USB MUX: Generalize mux ACK

Currently, only the virtual mux driver uses the mux ACK feature, but the
actual wait for the host command ACK is a part of the usb_mux general
code.  Generalize this mux ACK wait so it's available if needed in the
future for more muxes.

Additionally, moving this wait out of the mux set will allow us to lock
the muxes intelligently between tasks, without keeping the muxes locked
during the inactive ACK wait.

BRANCH=None
BUG=b:172222942,b:186777984
TEST=tast typec.Mode*.manual on voxel

Signed-off-by: Diana Z <dzigterman@chromium.org>
Change-Id: I61a043425a482cc6f3170548c888d91ec20c2a82
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/3078411
Reviewed-by: Keith Short <keithshort@chromium.org>

38 files changed, 248 insertions, 66 deletions

diff --git a/baseboard/guybrush/baseboard.c b/baseboard/guybrush/baseboard.c
index 4383598ff6..675350a022 100644
--- a/baseboard/guybrush/baseboard.c
+++ b/baseboard/guybrush/baseboard.c
@@ -384,7 +384,7 @@ BUILD_ASSERT(ARRAY_SIZE(pi3usb9201_bc12_chips) == USBC_PORT_COUNT);
  * not needed as well. usb_mux.c can handle the situation
  * properly.
  */
-static int fsusb42umx_set_mux(const struct usb_mux*, mux_state_t);
+static int fsusb42umx_set_mux(const struct usb_mux*, mux_state_t, bool *);
 struct usb_mux_driver usbc0_sbu_mux_driver = {
 	.set = fsusb42umx_set_mux,
 };
@@ -538,8 +538,12 @@ BUILD_ASSERT(ARRAY_SIZE(fans) == FAN_CH_COUNT);
  * chip and it needs a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/baseboard/mancomb/baseboard.c b/baseboard/mancomb/baseboard.c
index 041b2512dd..b79cf3a3de 100644
--- a/baseboard/mancomb/baseboard.c
+++ b/baseboard/mancomb/baseboard.c
@@ -341,7 +341,7 @@ BUILD_ASSERT(ARRAY_SIZE(pi3usb9201_bc12_chips) == USBC_PORT_COUNT);
  * not needed as well. usb_mux.c can handle the situation
  * properly.
  */
-static int fsusb42umx_set_mux(const struct usb_mux*, mux_state_t);
+static int fsusb42umx_set_mux(const struct usb_mux*, mux_state_t, bool *);
 const struct usb_mux_driver usbc_sbu_mux_driver = {
 	.set = fsusb42umx_set_mux,
 };
@@ -447,10 +447,14 @@ BUILD_ASSERT(ARRAY_SIZE(fans) == FAN_CH_COUNT);
  * chip and it needs a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
 	bool inverted = mux_state & USB_PD_MUX_POLARITY_INVERTED;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (me->usb_port == USBC_PORT_C0)
 		RETURN_ERROR(ioex_set_level(IOEX_USB_C0_SBU_FLIP, inverted));
 	else if (me->usb_port == USBC_PORT_C1)
diff --git a/board/berknip/board.c b/board/berknip/board.c
index e883300f1c..ca113980f6 100644
--- a/board/berknip/board.c
+++ b/board/berknip/board.c
@@ -121,8 +121,12 @@ DECLARE_HOOK(HOOK_CHIPSET_SUSPEND, board_chipset_suspend, HOOK_PRIO_DEFAULT);
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			     bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 0);
 	else
diff --git a/board/dalboz/board.c b/board/dalboz/board.c
index f049ef55d4..e4ee1bfe82 100644
--- a/board/dalboz/board.c
+++ b/board/dalboz/board.c
@@ -250,8 +250,12 @@ static int board_ps8743_mux_set(const struct usb_mux *me,
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/board/dirinboz/board.c b/board/dirinboz/board.c
index 36d3d13a04..78fd961573 100644
--- a/board/dirinboz/board.c
+++ b/board/dirinboz/board.c
@@ -171,8 +171,12 @@ DECLARE_HOOK(HOOK_CHIPSET_SUSPEND, retimers_off, HOOK_PRIO_DEFAULT);
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			     bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/board/ezkinil/board.c b/board/ezkinil/board.c
index 074a7694b0..f43cb116b0 100644
--- a/board/ezkinil/board.c
+++ b/board/ezkinil/board.c
@@ -276,8 +276,12 @@ const struct pi3hdx1204_tuning pi3hdx1204_tuning = {
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/board/gumboz/board.c b/board/gumboz/board.c
index ee255ad47b..4428819415 100644
--- a/board/gumboz/board.c
+++ b/board/gumboz/board.c
@@ -175,8 +175,12 @@ DECLARE_HOOK(HOOK_CHIPSET_SUSPEND, retimers_off, HOOK_PRIO_DEFAULT);
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int pi3usb221_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			     bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/board/pdeval-stm32f072/usb_pd_policy.c b/board/pdeval-stm32f072/usb_pd_policy.c
index d5a40fa6e8..57ef5a5826 100644
--- a/board/pdeval-stm32f072/usb_pd_policy.c
+++ b/board/pdeval-stm32f072/usb_pd_policy.c
@@ -204,6 +204,7 @@ __override int svdm_dp_config(int port, uint32_t *payload)
 {
 	int opos = pd_alt_mode(port, TCPC_TX_SOP, USB_SID_DISPLAYPORT);
 	int pin_mode = pd_dfp_dp_get_pin_mode(port, dp_status[port]);
+	bool unused;
 #if defined(CONFIG_USB_PD_TCPM_MUX) && defined(CONFIG_USB_PD_TCPM_ANX7447)
 	const struct usb_mux *mux = &usb_muxes[port];
 #endif
@@ -224,13 +225,17 @@ __override int svdm_dp_config(int port, uint32_t *payload)
 	case MODE_DP_PIN_C:
 	case MODE_DP_PIN_E:
 		mux_state |= USB_PD_MUX_DP_ENABLED;
-		mux->driver->set(mux, mux_state);
+		/*
+		 * Note: Direct mux driver calls are deprecated.  Calls
+		 * should go through the usb_mux APIs instead.
+		 */
+		mux->driver->set(mux, mux_state, &unused);
 		break;
 	case MODE_DP_PIN_B:
 	case MODE_DP_PIN_D:
 	case MODE_DP_PIN_F:
 		mux_state |= USB_PD_MUX_DOCK;
-		mux->driver->set(mux, mux_state);
+		mux->driver->set(mux, mux_state, &unused);
 		break;
 	}
 #endif
diff --git a/board/shuboz/board.c b/board/shuboz/board.c
index f5483841dc..84e2249e44 100644
--- a/board/shuboz/board.c
+++ b/board/shuboz/board.c
@@ -253,8 +253,12 @@ static int board_ps8743_mux_set(const struct usb_mux *me,
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
diff --git a/board/vilboz/board.c b/board/vilboz/board.c
index d280081e82..a39d46cbcd 100644
--- a/board/vilboz/board.c
+++ b/board/vilboz/board.c
@@ -156,8 +156,12 @@ unsigned int motion_sensor_count = ARRAY_SIZE(motion_sensors);
  * chip and it need a board specific driver.
  * Overall, it will use chained mux framework.
  */
-static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int fsusb42umx_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		ioex_set_level(IOEX_USB_C0_SBU_FLIP, 1);
 	else
@@ -429,8 +433,13 @@ static void lte_usb3_mux_init(void)
 			.i2c_addr_flags = AMD_FP5_MUX_I2C_ADDR_FLAGS,
 			.driver = &amd_fp5_usb_mux_driver,
 		};
+		bool unused;
+		/*
+		 * Note: Direct mux driver calls are deprecated.  Calls
+		 * should go through the usb_mux APIs instead.
+		 */
 		/* steer the mux to connect the USB 3 superspeed pairs */
-		usb_c1.driver->set(&usb_c1, USB_PD_MUX_USB_ENABLED);
+		usb_c1.driver->set(&usb_c1, USB_PD_MUX_USB_ENABLED, &unused);
 	}
 }
 DECLARE_HOOK(HOOK_CHIPSET_RESUME, lte_usb3_mux_init, HOOK_PRIO_DEFAULT);
diff --git a/common/mock/usb_mux_mock.c b/common/mock/usb_mux_mock.c
index e04c69f04b..f2db5cf8bd 100644
--- a/common/mock/usb_mux_mock.c
+++ b/common/mock/usb_mux_mock.c
@@ -31,8 +31,12 @@ static int mock_init(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int mock_set(const struct usb_mux *me, mux_state_t mux_state)
+static int mock_set(const struct usb_mux *me, mux_state_t mux_state,
+		    bool *ack_required)
 {
+	/* Mock does not use host command ACKs */
+	*ack_required = false;
+
 	mock_usb_mux.state = mux_state;
 	++mock_usb_mux.num_set_calls;
 	ccprints("[MUX] Set to 0x%02x", mux_state);
diff --git a/driver/retimer/bb_retimer.c b/driver/retimer/bb_retimer.c
index 919e278a8b..43c67a95e1 100644
--- a/driver/retimer/bb_retimer.c
+++ b/driver/retimer/bb_retimer.c
@@ -379,12 +379,16 @@ __overridable int bb_retimer_reset(const struct usb_mux *me)
 /**
  * Driver interface functions
  */
-static int retimer_set_state(const struct usb_mux *me, mux_state_t mux_state)
+static int retimer_set_state(const struct usb_mux *me, mux_state_t mux_state,
+			     bool *ack_required)
 {
 	uint32_t set_retimer_con = 0;
 	uint8_t dp_pin_mode;
 	int port = me->usb_port;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/*
 	 * Bit 0: DATA_CONNECTION_PRESENT
 	 * 0 - No connection present
diff --git a/driver/retimer/kb800x.c b/driver/retimer/kb800x.c
index 3fbb52cbc1..ad615a3532 100644
--- a/driver/retimer/kb800x.c
+++ b/driver/retimer/kb800x.c
@@ -376,10 +376,14 @@ static int kb800x_cio_init(const struct usb_mux *me, mux_state_t mux_state)
 	return kb800x_write(me, KB800X_REG_ORIENTATION, orientation);
 }
 
-static int kb800x_set_state(const struct usb_mux *me, mux_state_t mux_state)
+static int kb800x_set_state(const struct usb_mux *me, mux_state_t mux_state,
+			    bool *ack_required)
 {
 	int rv;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	cached_mux_state[me->usb_port] = mux_state;
 	rv = kb800x_write(me, KB800X_REG_RESET, KB800X_RESET_MASK);
 	if (rv)
diff --git a/driver/retimer/nb7v904m.c b/driver/retimer/nb7v904m.c
index 18c7809e71..94e96230b2 100644
--- a/driver/retimer/nb7v904m.c
+++ b/driver/retimer/nb7v904m.c
@@ -147,12 +147,16 @@ static int nb7v904m_init(const struct usb_mux *me)
 	return rv;
 }
 
-static int nb7v904m_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int nb7v904m_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			    bool *ack_required)
 {
 	int rv = EC_SUCCESS;
 	int regval;
 	int flipped = !!(mux_state & USB_PD_MUX_POLARITY_INVERTED);
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/* Turn off redriver if it's not needed at all. */
 	if (mux_state == USB_PD_MUX_NONE)
 		return nb7v904m_enter_low_power_mode(me);
diff --git a/driver/retimer/pi3dpx1207.c b/driver/retimer/pi3dpx1207.c
index 9254c3d4bb..8829c508a1 100644
--- a/driver/retimer/pi3dpx1207.c
+++ b/driver/retimer/pi3dpx1207.c
@@ -90,7 +90,8 @@ static int pi3dpx1207_enter_low_power_mode(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int pi3dpx1207_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int pi3dpx1207_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			      bool *ack_required)
 {
 	int rv = EC_SUCCESS;
 	uint8_t mode_val = PI3DPX1207_MODE_WATCHDOG_EN;
@@ -98,6 +99,9 @@ static int pi3dpx1207_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 	const int gpio_enable = pi3dpx1207_controls[port].enable_gpio;
 	const int gpio_dp_enable = pi3dpx1207_controls[port].dp_enable_gpio;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/* USB */
 	if (mux_state & USB_PD_MUX_USB_ENABLED) {
 		gpio_or_ioex_set_level(gpio_enable, 1);
diff --git a/driver/retimer/ps8802.c b/driver/retimer/ps8802.c
index 7406d8c54e..9738123ace 100644
--- a/driver/retimer/ps8802.c
+++ b/driver/retimer/ps8802.c
@@ -215,11 +215,15 @@ static int ps8802_init(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int ps8802_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8802_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	int val;
 	int rv;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return (mux_state == USB_PD_MUX_NONE) ? EC_SUCCESS
 						      : EC_ERROR_NOT_POWERED;
diff --git a/driver/retimer/ps8818.c b/driver/retimer/ps8818.c
index 65abcddd23..2f8e353099 100644
--- a/driver/retimer/ps8818.c
+++ b/driver/retimer/ps8818.c
@@ -98,11 +98,15 @@ int ps8818_i2c_field_update8(const struct usb_mux *me, int page, int offset,
 	return rv;
 }
 
-static int ps8818_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8818_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	int rv;
 	int val = 0;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return (mux_state == USB_PD_MUX_NONE) ? EC_SUCCESS
 						      : EC_ERROR_NOT_POWERED;
diff --git a/driver/retimer/tusb544.c b/driver/retimer/tusb544.c
index e76e5bd976..9de543fd42 100644
--- a/driver/retimer/tusb544.c
+++ b/driver/retimer/tusb544.c
@@ -58,11 +58,15 @@ static int tusb544_init(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int tusb544_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int tusb544_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int reg;
 	int rv;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state == USB_PD_MUX_NONE)
 		return tusb544_enter_low_power_mode(me);
 
diff --git a/driver/tcpm/anx7447.c b/driver/tcpm/anx7447.c
index 9458ee8c94..0367d4a78c 100644
--- a/driver/tcpm/anx7447.c
+++ b/driver/tcpm/anx7447.c
@@ -38,7 +38,8 @@ struct anx_usb_mux {
 	int state;
 };
 
-static int anx7447_mux_set(const struct usb_mux *me, mux_state_t mux_state);
+static int anx7447_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required);
 
 static struct anx_state anx[CONFIG_USB_PD_PORT_MAX_COUNT];
 static struct anx_usb_mux mux[CONFIG_USB_PD_PORT_MAX_COUNT];
@@ -290,6 +291,7 @@ static int anx7447_init(int port)
 {
 	int rv, reg, i;
 	const struct usb_mux *me = &usb_muxes[port];
+	bool unused;
 
 	ASSERT(port < CONFIG_USB_PD_PORT_MAX_COUNT);
 
@@ -388,9 +390,13 @@ static int anx7447_init(int port)
 	while ((me != NULL) && (me->driver != &anx7447_usb_mux_driver))
 		me = me->next_mux;
 
+	/*
+	 * Note that bypassing the usb_mux API is okay for internal driver calls
+	 * since the task calling init already holds this port's mux lock.
+	 */
 	if (me != NULL &&
 	    !(me->flags & USB_MUX_FLAG_NOT_TCPC))
-		rv = anx7447_mux_set(me, USB_PD_MUX_NONE);
+		rv = anx7447_mux_set(me, USB_PD_MUX_NONE, &unused);
 #endif /* CONFIG_USB_PD_TCPM_MUX */
 
 	return rv;
@@ -496,6 +502,7 @@ void anx7447_tcpc_clear_hpd_status(int port)
 static int anx7447_mux_init(const struct usb_mux *me)
 {
 	int port = me->usb_port;
+	bool unused;
 
 	ASSERT(port < CONFIG_USB_PD_PORT_MAX_COUNT);
 
@@ -511,7 +518,7 @@ static int anx7447_mux_init(const struct usb_mux *me)
 	 * USB_PD_MUX_DP_ENABLED) when reinitialized, we need to force
 	 * initialize it to USB_PD_MUX_NONE
 	 */
-	return anx7447_mux_set(me, USB_PD_MUX_NONE);
+	return anx7447_mux_set(me, USB_PD_MUX_NONE, &unused);
 }
 
 #ifdef CONFIG_USB_PD_TCPM_ANX7447_AUX_PU_PD
@@ -559,7 +566,8 @@ static inline void anx7447_configure_aux_src(const struct usb_mux *me,
  *
  * a2, a3, a10, a11, b2, b3, b10, b11 are pins on the USB-C connector.
  */
-static int anx7447_mux_set(const struct usb_mux *me, mux_state_t mux_state)
+static int anx7447_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int cc_direction;
 	mux_state_t mux_type;
@@ -567,6 +575,9 @@ static int anx7447_mux_set(const struct usb_mux *me, mux_state_t mux_state)
 	int rv;
 	int port = me->usb_port;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	cc_direction = mux_state & USB_PD_MUX_POLARITY_INVERTED;
 	mux_type = mux_state & USB_PD_MUX_DOCK;
 	CPRINTS("C%d mux_state = 0x%x, mux_type = 0x%x",
diff --git a/driver/tcpm/anx74xx.c b/driver/tcpm/anx74xx.c
index f45f7cd25e..8916fac17d 100644
--- a/driver/tcpm/anx74xx.c
+++ b/driver/tcpm/anx74xx.c
@@ -385,13 +385,17 @@ static int anx74xx_mux_aux_to_sbu(int port, int polarity, int enabled)
 }
 
 static int anx74xx_tcpm_mux_set(const struct usb_mux *me,
-				mux_state_t mux_state)
+				mux_state_t mux_state,
+				bool *ack_required)
 {
 	int ctrl5;
 	int ctrl1 = 0;
 	int rv;
 	int port = me->usb_port;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (!(mux_state & ~USB_PD_MUX_POLARITY_INVERTED)) {
 		anx[port].mux_state = mux_state;
 		return anx74xx_tcpm_mux_exit(port);
@@ -776,6 +780,7 @@ static int anx74xx_tcpm_set_polarity(int port, enum tcpc_cc_polarity polarity)
 {
 	int reg, mux_state, rv = EC_SUCCESS;
 	const struct usb_mux *me = &usb_muxes[port];
+	bool unused;
 
 	rv |= tcpc_read(port, ANX74XX_REG_CC_SOFTWARE_CTRL, &reg);
 	if (polarity_rm_dts(polarity)) /* Inform ANX to use CC2 */
@@ -791,7 +796,7 @@ static int anx74xx_tcpm_set_polarity(int port, enum tcpc_cc_polarity polarity)
 	mux_state = anx[port].mux_state & ~USB_PD_MUX_POLARITY_INVERTED;
 	if (polarity_rm_dts(polarity))
 		mux_state |= USB_PD_MUX_POLARITY_INVERTED;
-	anx74xx_tcpm_mux_set(me, mux_state);
+	anx74xx_tcpm_mux_set(me, mux_state, &unused);
 #endif
 	return rv;
 }
diff --git a/driver/tcpm/anx7688.c b/driver/tcpm/anx7688.c
index 9a2ed85de1..5e37352bc5 100644
--- a/driver/tcpm/anx7688.c
+++ b/driver/tcpm/anx7688.c
@@ -144,11 +144,15 @@ static void anx7688_tcpc_alert(int port)
 		anx7688_update_hpd_enable(port);
 }
 
-static int anx7688_mux_set(const struct usb_mux *me, mux_state_t mux_state)
+static int anx7688_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int reg = 0;
 	int rv, polarity;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	rv = mux_read(me, TCPC_REG_CONFIG_STD_OUTPUT, &reg);
 	if (rv != EC_SUCCESS)
 		return rv;
diff --git a/driver/tcpm/ps8xxx.c b/driver/tcpm/ps8xxx.c
index 58fb109b3e..ce0a4a9c6a 100644
--- a/driver/tcpm/ps8xxx.c
+++ b/driver/tcpm/ps8xxx.c
@@ -886,9 +886,9 @@ void ps8xxx_wake_from_standby(const struct usb_mux *me)
 	msleep(10);
 }
 
-static int ps8xxx_mux_set(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8xxx_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
-
 	if (product_id[me->usb_port] == PS8751_PRODUCT_ID &&
 	    me->flags & USB_MUX_FLAG_NOT_TCPC) {
 		ps8xxx_wake_from_standby(me);
@@ -906,7 +906,7 @@ static int ps8xxx_mux_set(const struct usb_mux *me, mux_state_t mux_state)
 							       TYPEC_CC_RD)));
 	}
 
-	return tcpci_tcpm_mux_set(me, mux_state);
+	return tcpci_tcpm_mux_set(me, mux_state, ack_required);
 }
 
 static int ps8xxx_mux_get(const struct usb_mux *me, mux_state_t *mux_state)
diff --git a/driver/tcpm/tcpci.c b/driver/tcpm/tcpci.c
index 3bd5d1b5d0..d113d876ea 100644
--- a/driver/tcpm/tcpci.c
+++ b/driver/tcpm/tcpci.c
@@ -1515,11 +1515,15 @@ int tcpci_tcpm_mux_enter_low_power(const struct usb_mux *me)
 	return mux_write(me, TCPC_REG_COMMAND, TCPC_REG_COMMAND_I2CIDLE);
 }
 
-int tcpci_tcpm_mux_set(const struct usb_mux *me, mux_state_t mux_state)
+int tcpci_tcpm_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+		       bool *ack_required)
 {
 	int rv;
 	int reg = 0;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/* Parameter is port only */
 	rv = mux_read(me, TCPC_REG_CONFIG_STD_OUTPUT, &reg);
 	if (rv != EC_SUCCESS)
diff --git a/driver/usb_mux/amd_fp5.c b/driver/usb_mux/amd_fp5.c
index 92c7e682a4..b77edf2826 100644
--- a/driver/usb_mux/amd_fp5.c
+++ b/driver/usb_mux/amd_fp5.c
@@ -42,10 +42,14 @@ static int amd_fp5_init(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int amd_fp5_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int amd_fp5_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	uint8_t val = 0;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	saved_mux_state[me->usb_port] = mux_state;
 
 	/*
@@ -128,9 +132,11 @@ static void amd_fp5_chipset_reset_delay(void)
 {
 	struct usb_mux *me;
 	int rv;
+	bool unused;
 
 	while (queue_remove_unit(&chipset_reset_queue, &me)) {
-		rv = amd_fp5_set_mux(me, saved_mux_state[me->usb_port]);
+		rv = amd_fp5_set_mux(me, saved_mux_state[me->usb_port],
+				     &unused);
 		if (rv)
 			ccprints("C%d restore mux rv:%d", me->usb_port, rv);
 	}
diff --git a/driver/usb_mux/amd_fp6.c b/driver/usb_mux/amd_fp6.c
index 88f6b5aff5..b2d5ae1fb4 100644
--- a/driver/usb_mux/amd_fp6.c
+++ b/driver/usb_mux/amd_fp6.c
@@ -134,10 +134,14 @@ static void amd_fp6_set_mux_retry(void)
 }
 
 
-static int amd_fp6_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int amd_fp6_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	uint8_t val;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state == USB_PD_MUX_NONE)
 		/*
 		 * LOW_POWER must be set when connection mode is
diff --git a/driver/usb_mux/anx3443.c b/driver/usb_mux/anx3443.c
index dc7fd3a3c0..2e57f4d30c 100644
--- a/driver/usb_mux/anx3443.c
+++ b/driver/usb_mux/anx3443.c
@@ -76,10 +76,14 @@ static int anx3443_wake_up(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int anx3443_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int anx3443_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int reg;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/* Mux is not powered in Z1 */
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return (mux_state == USB_PD_MUX_NONE) ? EC_SUCCESS
@@ -129,6 +133,7 @@ static int anx3443_get_mux(const struct usb_mux *me, mux_state_t *mux_state)
 static int anx3443_init(const struct usb_mux *me)
 {
 	uint64_t now;
+	bool unused;
 
 	/*
 	 * ANX3443 requires 30ms to power on. EC and ANX3443 are on the same
@@ -140,8 +145,12 @@ static int anx3443_init(const struct usb_mux *me)
 
 	RETURN_ERROR(anx3443_wake_up(me));
 
+	/*
+	 * Note that bypassing the usb_mux API is okay for internal driver calls
+	 * since the task calling init already holds this port's mux lock.
+	 */
 	/* Default to USB mode */
-	RETURN_ERROR(anx3443_set_mux(me, USB_PD_MUX_USB_ENABLED));
+	RETURN_ERROR(anx3443_set_mux(me, USB_PD_MUX_USB_ENABLED, &unused));
 
 	return EC_SUCCESS;
 }
diff --git a/driver/usb_mux/anx7440.c b/driver/usb_mux/anx7440.c
index df3a9c2682..89e593217d 100644
--- a/driver/usb_mux/anx7440.c
+++ b/driver/usb_mux/anx7440.c
@@ -61,10 +61,14 @@ static int anx7440_init(const struct usb_mux *me)
 }
 
 /* Writes control register to set switch mode */
-static int anx7440_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int anx7440_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int reg, res;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	res = anx7440_read(me, ANX7440_REG_CHIP_CTRL, &reg);
 	if (res)
 		return res;
diff --git a/driver/usb_mux/anx7451.c b/driver/usb_mux/anx7451.c
index ccf09f61be..52318eee29 100644
--- a/driver/usb_mux/anx7451.c
+++ b/driver/usb_mux/anx7451.c
@@ -89,10 +89,14 @@ static int anx7451_wake_up(const struct usb_mux *me)
 	return EC_SUCCESS;
 }
 
-static int anx7451_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int anx7451_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int reg;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	/* Mux is not powered in Z1 */
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return (mux_state == USB_PD_MUX_NONE) ? EC_SUCCESS
diff --git a/driver/usb_mux/it5205.c b/driver/usb_mux/it5205.c
index f275568313..0cfecdeda0 100644
--- a/driver/usb_mux/it5205.c
+++ b/driver/usb_mux/it5205.c
@@ -91,10 +91,14 @@ enum ec_error_list it5205h_enable_csbu_switch(const struct usb_mux *me, bool en)
 }
 
 /* Writes control register to set switch mode */
-static int it5205_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int it5205_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	uint8_t reg;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	switch (mux_state & MUX_STATE_DP_USB_MASK) {
 	case USB_PD_MUX_USB_ENABLED:
 		reg = IT5205_USB;
diff --git a/driver/usb_mux/pi3usb3x532.c b/driver/usb_mux/pi3usb3x532.c
index 7e74157d17..2435157967 100644
--- a/driver/usb_mux/pi3usb3x532.c
+++ b/driver/usb_mux/pi3usb3x532.c
@@ -84,10 +84,14 @@ static int pi3usb3x532_init(const struct usb_mux *me)
 
 /* Writes control register to set switch mode */
 static int pi3usb3x532_set_mux(const struct usb_mux *me,
-			       mux_state_t mux_state)
+			       mux_state_t mux_state,
+			       bool *ack_required)
 {
 	uint8_t reg = 0;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_USB_ENABLED)
 		reg |= PI3USB3X532_MODE_USB;
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
diff --git a/driver/usb_mux/ps8740.c b/driver/usb_mux/ps8740.c
index 4d03cb5cf3..618c74cd65 100644
--- a/driver/usb_mux/ps8740.c
+++ b/driver/usb_mux/ps8740.c
@@ -70,10 +70,14 @@ static int ps8740_init(const struct usb_mux *me)
 }
 
 /* Writes control register to set switch mode */
-static int ps8740_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8740_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	uint8_t reg = 0;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_USB_ENABLED)
 		reg |= PS8740_MODE_USB_ENABLED;
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
diff --git a/driver/usb_mux/ps8743.c b/driver/usb_mux/ps8743.c
index 68bb0428d3..f618bb009f 100644
--- a/driver/usb_mux/ps8743.c
+++ b/driver/usb_mux/ps8743.c
@@ -94,7 +94,8 @@ static int ps8743_init(const struct usb_mux *me)
 }
 
 /* Writes control register to set switch mode */
-static int ps8743_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8743_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	/*
 	 * For CE_DP, CE_USB, and FLIP, disable pin control and enable I2C
@@ -105,6 +106,9 @@ static int ps8743_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 		       PS8743_MODE_USB_REG_CONTROL |
 		       PS8743_MODE_FLIP_REG_CONTROL);
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_USB_ENABLED)
 		reg |= PS8743_MODE_USB_ENABLE;
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
diff --git a/driver/usb_mux/ps8822.c b/driver/usb_mux/ps8822.c
index f7a60df6fa..7f25db37f4 100644
--- a/driver/usb_mux/ps8822.c
+++ b/driver/usb_mux/ps8822.c
@@ -79,11 +79,15 @@ static int ps8822_init(const struct usb_mux *me)
 }
 
 /* Writes control register to set switch mode */
-static int ps8822_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int ps8822_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			  bool *ack_required)
 {
 	int reg;
 	int rv;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	rv = ps8822_read(me, PS8822_REG_PAGE0, PS8822_REG_MODE, &reg);
 	if (rv)
 		return rv;
diff --git a/driver/usb_mux/tusb1064.c b/driver/usb_mux/tusb1064.c
index 6bbf323b3c..1c0f0e4701 100644
--- a/driver/usb_mux/tusb1064.c
+++ b/driver/usb_mux/tusb1064.c
@@ -32,10 +32,14 @@ static int tusb1064_write(const struct usb_mux *me, uint8_t reg, uint8_t val)
 }
 
 /* Writes control register to set switch mode */
-static int tusb1064_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int tusb1064_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			    bool *ack_required)
 {
 	int reg = REG_GENERAL_STATIC_BITS;
 
+	/* This driver does not use host command ACKs */
+	*ack_required = false;
+
 	if (mux_state & USB_PD_MUX_USB_ENABLED)
 		reg |= REG_GENERAL_CTLSEL_USB3;
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
@@ -72,6 +76,7 @@ static int tusb1064_init(const struct usb_mux *me)
 {
 	int res;
 	uint8_t reg;
+	bool unused;
 
 	/* Default to "Floating Pin" DP Equalization */
 	reg = TUSB1064_DP1EQ(TUSB1064_DP_EQ_RX_10_0_DB) |
@@ -86,8 +91,12 @@ static int tusb1064_init(const struct usb_mux *me)
 	if (res)
 		return res;
 
+	/*
+	 * Note that bypassing the usb_mux API is okay for internal driver calls
+	 * since the task calling init already holds this port's mux lock.
+	 */
 	/* Disconnect USB3.1 and DP */
-	res = tusb1064_set_mux(me, USB_PD_MUX_NONE);
+	res = tusb1064_set_mux(me, USB_PD_MUX_NONE, &unused);
 	if (res)
 		return res;
 
diff --git a/driver/usb_mux/usb_mux.c b/driver/usb_mux/usb_mux.c
index f3418556d0..cf2e660a4f 100644
--- a/driver/usb_mux/usb_mux.c
+++ b/driver/usb_mux/usb_mux.c
@@ -12,6 +12,7 @@
 #include "hooks.h"
 #include "host_command.h"
 #include "task.h"
+#include "timer.h"
 #include "usb_mux.h"
 #include "usbc_ppc.h"
 #include "util.h"
@@ -73,6 +74,7 @@ static int configure_mux(int port,
 	     mux_ptr = mux_ptr->next_mux) {
 		mux_state_t lcl_state;
 		const struct usb_mux_driver *drv = mux_ptr->driver;
+		bool ack_required = false;
 
 		switch (config) {
 		case USB_MUX_INIT:
@@ -107,7 +109,8 @@ static int configure_mux(int port,
 				lcl_state &= ~USB_PD_MUX_POLARITY_INVERTED;
 
 			if (drv && drv->set) {
-				rv = drv->set(mux_ptr, lcl_state);
+				rv = drv->set(mux_ptr, lcl_state,
+					      &ack_required);
 				if (rv)
 					break;
 			}
@@ -133,6 +136,21 @@ static int configure_mux(int port,
 			}
 			break;
 		}
+
+		if (ack_required) {
+			/* This should only be called from the PD task */
+			assert(port == TASK_ID_TO_PD_PORT(task_get_current()));
+
+			/*
+			 * Note: This task event could be generalized for more
+			 * purposes beyond host command ACKs.  For now, these
+			 * wait times are tuned for the purposes of the TCSS
+			 * mux, but could be made configurable for other
+			 * purposes.
+			 */
+			task_wait_event_mask(PD_EVENT_AP_MUX_DONE, 100*MSEC);
+			usleep(12.5 * MSEC);
+		}
 	}
 
 	if (rv)
diff --git a/driver/usb_mux/virtual.c b/driver/usb_mux/virtual.c
index 9fbcbd2643..72020dda03 100644
--- a/driver/usb_mux/virtual.c
+++ b/driver/usb_mux/virtual.c
@@ -26,12 +26,19 @@
 
 static mux_state_t virtual_mux_state[CONFIG_USB_PD_PORT_MAX_COUNT];
 
-static inline void virtual_mux_update_state(int port, mux_state_t mux_state)
+static inline void virtual_mux_update_state(int port, mux_state_t mux_state,
+					    bool *ack_required)
 {
 	mux_state_t previous_mux_state = virtual_mux_state[port];
 
 	virtual_mux_state[port] = mux_state;
 
+	/*
+	 * Initialize ack_required to false to start, and set on necessary
+	 * conditions
+	 */
+	*ack_required = false;
+
 	if (!IS_ENABLED(CONFIG_HOSTCMD_EVENTS))
 		return;
 
@@ -48,9 +55,6 @@ static inline void virtual_mux_update_state(int port, mux_state_t mux_state)
 	 * TCSS Mux to allow better synchronization between them and thereby
 	 * remain in the same state for achieving proper safe state
 	 * terminations.
-	 *
-	 * Note: While the EC waits for the ACK, the value of usb_mux_get
-	 * won't match the most recently set value with usb_mux_set.
 	 */
 
 	/* TODO(b/186777984): Wait for an ACK for all mux state change */
@@ -60,13 +64,8 @@ static inline void virtual_mux_update_state(int port, mux_state_t mux_state)
 	   ((previous_mux_state & USB_PD_MUX_SAFE_MODE) &&
 	    !(mux_state & USB_PD_MUX_SAFE_MODE)) ||
 	   ((previous_mux_state != USB_PD_MUX_NONE) &&
-	    (mux_state == USB_PD_MUX_NONE))) {
-		/* This should only be called from the PD task */
-		assert(port == TASK_ID_TO_PD_PORT(task_get_current()));
-
-		task_wait_event_mask(PD_EVENT_AP_MUX_DONE, 100*MSEC);
-		usleep(12.5 * MSEC);
-	}
+	    (mux_state == USB_PD_MUX_NONE)))
+		*ack_required = true;
 }
 
 static int virtual_init(const struct usb_mux *me)
@@ -78,7 +77,8 @@ static int virtual_init(const struct usb_mux *me)
  * Set the state of our 'virtual' mux. The EC does not actually control this
  * mux, so update the desired state, then notify the host of the update.
  */
-static int virtual_set_mux(const struct usb_mux *me, mux_state_t mux_state)
+static int virtual_set_mux(const struct usb_mux *me, mux_state_t mux_state,
+			   bool *ack_required)
 {
 	int port = me->usb_port;
 
@@ -86,7 +86,7 @@ static int virtual_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 	mux_state_t new_mux_state = (mux_state & ~USB_PD_MUX_HPD_STATE) |
 		(virtual_mux_state[port] & USB_PD_MUX_HPD_STATE);
 
-	virtual_mux_update_state(port, new_mux_state);
+	virtual_mux_update_state(port, new_mux_state, ack_required);
 
 	return EC_SUCCESS;
 }
@@ -108,13 +108,15 @@ static int virtual_get_mux(const struct usb_mux *me, mux_state_t *mux_state)
 void virtual_hpd_update(const struct usb_mux *me, int hpd_lvl, int hpd_irq)
 {
 	int port = me->usb_port;
+	bool unused;
 
 	/* Current HPD related mux status + existing USB & DP mux status */
 	mux_state_t new_mux_state = (hpd_lvl ? USB_PD_MUX_HPD_LVL : 0) |
 			(hpd_irq ? USB_PD_MUX_HPD_IRQ : 0) |
 			(virtual_mux_state[port] & USB_PD_MUX_USB_DP_STATE);
 
-	virtual_mux_update_state(port, new_mux_state);
+	/* HPD ACK isn't required for the EC to continue with its tasks */
+	virtual_mux_update_state(port, new_mux_state, &unused);
 }
 
 const struct usb_mux_driver virtual_usb_mux_driver = {
diff --git a/include/driver/tcpm/tcpci.h b/include/driver/tcpm/tcpci.h
index 669236bcb0..04002862df 100644
--- a/include/driver/tcpm/tcpci.h
+++ b/include/driver/tcpm/tcpci.h
@@ -282,7 +282,8 @@ void tcpci_tcpc_enable_auto_discharge_disconnect(int port, int enable);
 int tcpci_tcpc_debug_accessory(int port, bool enable);
 
 int tcpci_tcpm_mux_init(const struct usb_mux *me);
-int tcpci_tcpm_mux_set(const struct usb_mux *me, mux_state_t mux_state);
+int tcpci_tcpm_mux_set(const struct usb_mux *me, mux_state_t mux_state,
+		       bool *ack_required);
 int tcpci_tcpm_mux_get(const struct usb_mux *me, mux_state_t *mux_state);
 int tcpci_tcpm_mux_enter_low_power(const struct usb_mux *me);
 int tcpci_get_chip_info(int port, int live,
diff --git a/include/usb_mux.h b/include/usb_mux.h
index 61a2c5167b..548ee8940b 100644
--- a/include/usb_mux.h
+++ b/include/usb_mux.h
@@ -42,11 +42,14 @@ struct usb_mux_driver {
 	/**
 	 * Set USB mux state.
 	 *
-	 * @param me usb_mux
-	 * @param mux_state State to set mux to.
+	 * @param[in]  me usb_mux
+	 * @param[in]  mux_state State to set mux to.
+	 * @param[out] bool ack_required - indication of whether this mux needs
+	 * to wait on a host command ACK at the end of a set
 	 * @return EC_SUCCESS on success, non-zero error code on failure.
 	 */
-	int (*set)(const struct usb_mux *me, mux_state_t mux_state);
+	int (*set)(const struct usb_mux *me, mux_state_t mux_state,
+		   bool *ack_required);
 
 	/**
 	 * Get current state of USB mux.