thunderbolt: Add support for enhanced uni-directional TMU mode

This is new TMU mode introduced with the USB4 v2. This mode is simpler
than the existing ones and allows all CL states as well. Enable this for
all links where both side routers are v2 and keep the existing
functionality for the v1 and earlier links.

Currently only support the MedRes rate. We can add the HiFi rate later
too if it turns out to be useful.

Signed-off-by: Mika Westerberg <mika.westerberg@linux.intel.com>

1 files changed, 471 insertions, 124 deletions

diff --git a/drivers/thunderbolt/tmu.c b/drivers/thunderbolt/tmu.c
index c926fb71c43d..1269f417515b 100644
--- a/drivers/thunderbolt/tmu.c
+++ b/drivers/thunderbolt/tmu.c
@@ -11,23 +11,63 @@
 
 #include "tb.h"
 
+static const unsigned int tmu_rates[] = {
+	[TB_SWITCH_TMU_MODE_OFF] = 0,
+	[TB_SWITCH_TMU_MODE_LOWRES] = 1000,
+	[TB_SWITCH_TMU_MODE_HIFI_UNI] = 16,
+	[TB_SWITCH_TMU_MODE_HIFI_BI] = 16,
+	[TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI] = 16,
+};
+
+const struct {
+	unsigned int freq_meas_window;
+	unsigned int avg_const;
+	unsigned int delta_avg_const;
+	unsigned int repl_timeout;
+	unsigned int repl_threshold;
+	unsigned int repl_n;
+	unsigned int dirswitch_n;
+} tmu_params[] = {
+	[TB_SWITCH_TMU_MODE_OFF] = { },
+	[TB_SWITCH_TMU_MODE_LOWRES] = { 30, 4, },
+	[TB_SWITCH_TMU_MODE_HIFI_UNI] = { 800, 8, },
+	[TB_SWITCH_TMU_MODE_HIFI_BI] = { 800, 8, },
+	[TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI] = {
+		800, 4, 0, 3125, 25, 128, 255,
+	},
+};
+
+static const char *tmu_mode_name(enum tb_switch_tmu_mode mode)
+{
+	switch (mode) {
+	case TB_SWITCH_TMU_MODE_OFF:
+		return "off";
+	case TB_SWITCH_TMU_MODE_LOWRES:
+		return "uni-directional, LowRes";
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		return "uni-directional, HiFi";
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		return "bi-directional, HiFi";
+	case TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI:
+		return "enhanced uni-directional, MedRes";
+	default:
+		return "unknown";
+	}
+}
+
+static bool tb_switch_tmu_enhanced_is_supported(const struct tb_switch *sw)
+{
+	return usb4_switch_version(sw) > 1;
+}
+
 static int tb_switch_set_tmu_mode_params(struct tb_switch *sw,
-					 enum tb_switch_tmu_rate rate)
+					 enum tb_switch_tmu_mode mode)
 {
-	u32 freq_meas_wind[2] = { 30, 800 };
-	u32 avg_const[2] = { 4, 8 };
 	u32 freq, avg, val;
 	int ret;
 
-	if (rate == TB_SWITCH_TMU_RATE_NORMAL) {
-		freq = freq_meas_wind[0];
-		avg = avg_const[0];
-	} else if (rate == TB_SWITCH_TMU_RATE_HIFI) {
-		freq = freq_meas_wind[1];
-		avg = avg_const[1];
-	} else {
-		return 0;
-	}
+	freq = tmu_params[mode].freq_meas_window;
+	avg = tmu_params[mode].avg_const;
 
 	ret = tb_sw_read(sw, &val, TB_CFG_SWITCH,
 			 sw->tmu.cap + TMU_RTR_CS_0, 1);
@@ -56,37 +96,30 @@ static int tb_switch_set_tmu_mode_params(struct tb_switch *sw,
 		FIELD_PREP(TMU_RTR_CS_15_OFFSET_AVG_MASK, avg) |
 		FIELD_PREP(TMU_RTR_CS_15_ERROR_AVG_MASK, avg);
 
-	return tb_sw_write(sw, &val, TB_CFG_SWITCH,
-			   sw->tmu.cap + TMU_RTR_CS_15, 1);
-}
-
-static const char *tb_switch_tmu_mode_name(const struct tb_switch *sw)
-{
-	bool root_switch = !tb_route(sw);
+	ret = tb_sw_write(sw, &val, TB_CFG_SWITCH,
+			 sw->tmu.cap + TMU_RTR_CS_15, 1);
+	if (ret)
+		return ret;
 
-	switch (sw->tmu.rate) {
-	case TB_SWITCH_TMU_RATE_OFF:
-		return "off";
+	if (tb_switch_tmu_enhanced_is_supported(sw)) {
+		u32 delta_avg = tmu_params[mode].delta_avg_const;
 
-	case TB_SWITCH_TMU_RATE_HIFI:
-		/* Root switch does not have upstream directionality */
-		if (root_switch)
-			return "HiFi";
-		if (sw->tmu.unidirectional)
-			return "uni-directional, HiFi";
-		return "bi-directional, HiFi";
+		ret = tb_sw_read(sw, &val, TB_CFG_SWITCH,
+				 sw->tmu.cap + TMU_RTR_CS_18, 1);
+		if (ret)
+			return ret;
 
-	case TB_SWITCH_TMU_RATE_NORMAL:
-		if (root_switch)
-			return "normal";
-		return "uni-directional, normal";
+		val &= ~TMU_RTR_CS_18_DELTA_AVG_CONST_MASK;
+		val |= FIELD_PREP(TMU_RTR_CS_18_DELTA_AVG_CONST_MASK, delta_avg);
 
-	default:
-		return "unknown";
+		ret = tb_sw_write(sw, &val, TB_CFG_SWITCH,
+				  sw->tmu.cap + TMU_RTR_CS_18, 1);
 	}
+
+	return ret;
 }
 
-static bool tb_switch_tmu_ucap_supported(struct tb_switch *sw)
+static bool tb_switch_tmu_ucap_is_supported(struct tb_switch *sw)
 {
 	int ret;
 	u32 val;
@@ -182,6 +215,103 @@ static bool tb_port_tmu_is_unidirectional(struct tb_port *port)
 	return val & TMU_ADP_CS_3_UDM;
 }
 
+static bool tb_port_tmu_is_enhanced(struct tb_port *port)
+{
+	int ret;
+	u32 val;
+
+	ret = tb_port_read(port, &val, TB_CFG_PORT,
+			   port->cap_tmu + TMU_ADP_CS_8, 1);
+	if (ret)
+		return false;
+
+	return val & TMU_ADP_CS_8_EUDM;
+}
+
+/* Can be called to non-v2 lane adapters too */
+static int tb_port_tmu_enhanced_enable(struct tb_port *port, bool enable)
+{
+	int ret;
+	u32 val;
+
+	if (!tb_switch_tmu_enhanced_is_supported(port->sw))
+		return 0;
+
+	ret = tb_port_read(port, &val, TB_CFG_PORT,
+			   port->cap_tmu + TMU_ADP_CS_8, 1);
+	if (ret)
+		return ret;
+
+	if (enable)
+		val |= TMU_ADP_CS_8_EUDM;
+	else
+		val &= ~TMU_ADP_CS_8_EUDM;
+
+	return tb_port_write(port, &val, TB_CFG_PORT,
+			     port->cap_tmu + TMU_ADP_CS_8, 1);
+}
+
+static int tb_port_set_tmu_mode_params(struct tb_port *port,
+				       enum tb_switch_tmu_mode mode)
+{
+	u32 repl_timeout, repl_threshold, repl_n, dirswitch_n, val;
+	int ret;
+
+	repl_timeout = tmu_params[mode].repl_timeout;
+	repl_threshold = tmu_params[mode].repl_threshold;
+	repl_n = tmu_params[mode].repl_n;
+	dirswitch_n = tmu_params[mode].dirswitch_n;
+
+	ret = tb_port_read(port, &val, TB_CFG_PORT,
+			   port->cap_tmu + TMU_ADP_CS_8, 1);
+	if (ret)
+		return ret;
+
+	val &= ~TMU_ADP_CS_8_REPL_TIMEOUT_MASK;
+	val &= ~TMU_ADP_CS_8_REPL_THRESHOLD_MASK;
+	val |= FIELD_PREP(TMU_ADP_CS_8_REPL_TIMEOUT_MASK, repl_timeout);
+	val |= FIELD_PREP(TMU_ADP_CS_8_REPL_THRESHOLD_MASK, repl_threshold);
+
+	ret = tb_port_write(port, &val, TB_CFG_PORT,
+			    port->cap_tmu + TMU_ADP_CS_8, 1);
+	if (ret)
+		return ret;
+
+	ret = tb_port_read(port, &val, TB_CFG_PORT,
+			   port->cap_tmu + TMU_ADP_CS_9, 1);
+	if (ret)
+		return ret;
+
+	val &= ~TMU_ADP_CS_9_REPL_N_MASK;
+	val &= ~TMU_ADP_CS_9_DIRSWITCH_N_MASK;
+	val |= FIELD_PREP(TMU_ADP_CS_9_REPL_N_MASK, repl_n);
+	val |= FIELD_PREP(TMU_ADP_CS_9_DIRSWITCH_N_MASK, dirswitch_n);
+
+	return tb_port_write(port, &val, TB_CFG_PORT,
+			     port->cap_tmu + TMU_ADP_CS_9, 1);
+}
+
+/* Can be called to non-v2 lane adapters too */
+static int tb_port_tmu_rate_write(struct tb_port *port, int rate)
+{
+	int ret;
+	u32 val;
+
+	if (!tb_switch_tmu_enhanced_is_supported(port->sw))
+		return 0;
+
+	ret = tb_port_read(port, &val, TB_CFG_PORT,
+			   port->cap_tmu + TMU_ADP_CS_9, 1);
+	if (ret)
+		return ret;
+
+	val &= ~TMU_ADP_CS_9_ADP_TS_INTERVAL_MASK;
+	val |= FIELD_PREP(TMU_ADP_CS_9_ADP_TS_INTERVAL_MASK, rate);
+
+	return tb_port_write(port, &val, TB_CFG_PORT,
+			     port->cap_tmu + TMU_ADP_CS_9, 1);
+}
+
 static int tb_port_tmu_time_sync(struct tb_port *port, bool time_sync)
 {
 	u32 val = time_sync ? TMU_ADP_CS_6_DTS : 0;
@@ -224,6 +354,50 @@ static int tb_switch_tmu_set_time_disruption(struct tb_switch *sw, bool set)
 	return tb_sw_write(sw, &val, TB_CFG_SWITCH, offset, 1);
 }
 
+static int tmu_mode_init(struct tb_switch *sw)
+{
+	bool enhanced, ucap;
+	int ret, rate;
+
+	ucap = tb_switch_tmu_ucap_is_supported(sw);
+	if (ucap)
+		tb_sw_dbg(sw, "TMU: supports uni-directional mode\n");
+	enhanced = tb_switch_tmu_enhanced_is_supported(sw);
+	if (enhanced)
+		tb_sw_dbg(sw, "TMU: supports enhanced uni-directional mode\n");
+
+	ret = tb_switch_tmu_rate_read(sw);
+	if (ret < 0)
+		return ret;
+	rate = ret;
+
+	/* Off by default */
+	sw->tmu.mode = TB_SWITCH_TMU_MODE_OFF;
+
+	if (tb_route(sw)) {
+		struct tb_port *up = tb_upstream_port(sw);
+
+		if (enhanced && tb_port_tmu_is_enhanced(up)) {
+			sw->tmu.mode = TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI;
+		} else if (ucap && tb_port_tmu_is_unidirectional(up)) {
+			if (tmu_rates[TB_SWITCH_TMU_MODE_LOWRES] == rate)
+				sw->tmu.mode = TB_SWITCH_TMU_MODE_LOWRES;
+			else if (tmu_rates[TB_SWITCH_TMU_MODE_LOWRES] == rate)
+				sw->tmu.mode = TB_SWITCH_TMU_MODE_HIFI_UNI;
+		} else if (rate) {
+			sw->tmu.mode = TB_SWITCH_TMU_MODE_HIFI_BI;
+		}
+	} else if (rate) {
+		sw->tmu.mode = TB_SWITCH_TMU_MODE_HIFI_BI;
+	}
+
+	/* Update the initial request to match the current mode */
+	sw->tmu.mode_request = sw->tmu.mode;
+	sw->tmu.has_ucap = ucap;
+
+	return 0;
+}
+
 /**
  * tb_switch_tmu_init() - Initialize switch TMU structures
  * @sw: Switch to initialized
@@ -252,27 +426,11 @@ int tb_switch_tmu_init(struct tb_switch *sw)
 			port->cap_tmu = cap;
 	}
 
-	ret = tb_switch_tmu_rate_read(sw);
-	if (ret < 0)
+	ret = tmu_mode_init(sw);
+	if (ret)
 		return ret;
 
-	sw->tmu.rate = ret;
-
-	sw->tmu.has_ucap = tb_switch_tmu_ucap_supported(sw);
-	if (sw->tmu.has_ucap) {
-		tb_sw_dbg(sw, "TMU: supports uni-directional mode\n");
-
-		if (tb_route(sw)) {
-			struct tb_port *up = tb_upstream_port(sw);
-
-			sw->tmu.unidirectional =
-				tb_port_tmu_is_unidirectional(up);
-		}
-	} else {
-		sw->tmu.unidirectional = false;
-	}
-
-	tb_sw_dbg(sw, "TMU: current mode: %s\n", tb_switch_tmu_mode_name(sw));
+	tb_sw_dbg(sw, "TMU: current mode: %s\n", tmu_mode_name(sw->tmu.mode));
 	return 0;
 }
 
@@ -375,6 +533,23 @@ out:
 	return ret;
 }
 
+static int disable_enhanced(struct tb_port *up, struct tb_port *down)
+{
+	int ret;
+
+	/*
+	 * Router may already been disconnected so ignore errors on the
+	 * upstream port.
+	 */
+	tb_port_tmu_rate_write(up, 0);
+	tb_port_tmu_enhanced_enable(up, false);
+
+	ret = tb_port_tmu_rate_write(down, 0);
+	if (ret)
+		return ret;
+	return tb_port_tmu_enhanced_enable(down, false);
+}
+
 /**
  * tb_switch_tmu_disable() - Disable TMU of a switch
  * @sw: Switch whose TMU to disable
@@ -384,11 +559,10 @@ out:
 int tb_switch_tmu_disable(struct tb_switch *sw)
 {
 	/* Already disabled? */
-	if (sw->tmu.rate == TB_SWITCH_TMU_RATE_OFF)
+	if (sw->tmu.mode == TB_SWITCH_TMU_MODE_OFF)
 		return 0;
 
 	if (tb_route(sw)) {
-		bool unidirectional = sw->tmu.unidirectional;
 		struct tb_port *down, *up;
 		int ret;
 
@@ -405,33 +579,46 @@ int tb_switch_tmu_disable(struct tb_switch *sw)
 		 * uni-directional mode and we don't want to change it's TMU
 		 * mode.
 		 */
-		tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_OFF);
+		tb_switch_tmu_rate_write(sw, tmu_rates[TB_SWITCH_TMU_MODE_OFF]);
 
 		tb_port_tmu_time_sync_disable(up);
 		ret = tb_port_tmu_time_sync_disable(down);
 		if (ret)
 			return ret;
 
-		if (unidirectional) {
+		switch (sw->tmu.mode) {
+		case TB_SWITCH_TMU_MODE_LOWRES:
+		case TB_SWITCH_TMU_MODE_HIFI_UNI:
 			/* The switch may be unplugged so ignore any errors */
 			tb_port_tmu_unidirectional_disable(up);
 			ret = tb_port_tmu_unidirectional_disable(down);
 			if (ret)
 				return ret;
+			break;
+
+		case TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI:
+			ret = disable_enhanced(up, down);
+			if (ret)
+				return ret;
+			break;
+
+		default:
+			break;
 		}
 	} else {
-		tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_OFF);
+		tb_switch_tmu_rate_write(sw, tmu_rates[TB_SWITCH_TMU_MODE_OFF]);
 	}
 
-	sw->tmu.unidirectional = false;
-	sw->tmu.rate = TB_SWITCH_TMU_RATE_OFF;
+	sw->tmu.mode = TB_SWITCH_TMU_MODE_OFF;
 
 	tb_sw_dbg(sw, "TMU: disabled\n");
 	return 0;
 }
 
-static void tb_switch_tmu_off(struct tb_switch *sw, bool unidirectional)
+/* Called only when there is failure enabling requested mode */
+static void tb_switch_tmu_off(struct tb_switch *sw)
 {
+	unsigned int rate = tmu_rates[TB_SWITCH_TMU_MODE_OFF];
 	struct tb_port *down, *up;
 
 	down = tb_switch_downstream_port(sw);
@@ -445,20 +632,30 @@ static void tb_switch_tmu_off(struct tb_switch *sw, bool unidirectional)
 	 */
 	tb_port_tmu_time_sync_disable(down);
 	tb_port_tmu_time_sync_disable(up);
-	if (unidirectional)
-		tb_switch_tmu_rate_write(tb_switch_parent(sw),
-					 TB_SWITCH_TMU_RATE_OFF);
-	else
-		tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_OFF);
 
-	tb_switch_set_tmu_mode_params(sw, sw->tmu.rate);
+	switch (sw->tmu.mode_request) {
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		tb_switch_tmu_rate_write(tb_switch_parent(sw), rate);
+		break;
+	case TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI:
+		disable_enhanced(up, down);
+		break;
+	default:
+		break;
+	}
+
+	/* Always set the rate to 0 */
+	tb_switch_tmu_rate_write(sw, rate);
+
+	tb_switch_set_tmu_mode_params(sw, sw->tmu.mode);
 	tb_port_tmu_unidirectional_disable(down);
 	tb_port_tmu_unidirectional_disable(up);
 }
 
 /*
  * This function is called when the previous TMU mode was
- * TB_SWITCH_TMU_RATE_OFF.
+ * TB_SWITCH_TMU_MODE_OFF.
  */
 static int tb_switch_tmu_enable_bidirectional(struct tb_switch *sw)
 {
@@ -476,7 +673,7 @@ static int tb_switch_tmu_enable_bidirectional(struct tb_switch *sw)
 	if (ret)
 		goto out;
 
-	ret = tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_HIFI);
+	ret = tb_switch_tmu_rate_write(sw, tmu_rates[TB_SWITCH_TMU_MODE_HIFI_BI]);
 	if (ret)
 		goto out;
 
@@ -491,7 +688,7 @@ static int tb_switch_tmu_enable_bidirectional(struct tb_switch *sw)
 	return 0;
 
 out:
-	tb_switch_tmu_off(sw, false);
+	tb_switch_tmu_off(sw);
 	return ret;
 }
 
@@ -522,7 +719,7 @@ static int tb_switch_tmu_disable_objections(struct tb_switch *sw)
 
 /*
  * This function is called when the previous TMU mode was
- * TB_SWITCH_TMU_RATE_OFF.
+ * TB_SWITCH_TMU_MODE_OFF.
  */
 static int tb_switch_tmu_enable_unidirectional(struct tb_switch *sw)
 {
@@ -532,11 +729,11 @@ static int tb_switch_tmu_enable_unidirectional(struct tb_switch *sw)
 	up = tb_upstream_port(sw);
 	down = tb_switch_downstream_port(sw);
 	ret = tb_switch_tmu_rate_write(tb_switch_parent(sw),
-				       sw->tmu.rate_request);
+				       tmu_rates[sw->tmu.mode_request]);
 	if (ret)
 		return ret;
 
-	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.rate_request);
+	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.mode_request);
 	if (ret)
 		return ret;
 
@@ -559,12 +756,62 @@ static int tb_switch_tmu_enable_unidirectional(struct tb_switch *sw)
 	return 0;
 
 out:
-	tb_switch_tmu_off(sw, true);
+	tb_switch_tmu_off(sw);
+	return ret;
+}
+
+/*
+ * This function is called when the previous TMU mode was
+ * TB_SWITCH_TMU_RATE_OFF.
+ */
+static int tb_switch_tmu_enable_enhanced(struct tb_switch *sw)
+{
+	unsigned int rate = tmu_rates[sw->tmu.mode_request];
+	struct tb_port *up, *down;
+	int ret;
+
+	/* Router specific parameters first */
+	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.mode_request);
+	if (ret)
+		return ret;
+
+	up = tb_upstream_port(sw);
+	down = tb_switch_downstream_port(sw);
+
+	ret = tb_port_set_tmu_mode_params(up, sw->tmu.mode_request);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_rate_write(up, rate);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_enhanced_enable(up, true);
+	if (ret)
+		goto out;
+
+	ret = tb_port_set_tmu_mode_params(down, sw->tmu.mode_request);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_rate_write(down, rate);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_enhanced_enable(down, true);
+	if (ret)
+		goto out;
+
+	return 0;
+
+out:
+	tb_switch_tmu_off(sw);
 	return ret;
 }
 
 static void tb_switch_tmu_change_mode_prev(struct tb_switch *sw)
 {
+	unsigned int rate = tmu_rates[sw->tmu.mode];
 	struct tb_port *down, *up;
 
 	down = tb_switch_downstream_port(sw);
@@ -575,42 +822,97 @@ static void tb_switch_tmu_change_mode_prev(struct tb_switch *sw)
 	 * In case of additional failures in the functions below,
 	 * ignore them since the caller shall already report a failure.
 	 */
-	tb_port_tmu_set_unidirectional(down, sw->tmu.unidirectional);
-	if (sw->tmu.unidirectional_request)
-		tb_switch_tmu_rate_write(tb_switch_parent(sw), sw->tmu.rate);
-	else
-		tb_switch_tmu_rate_write(sw, sw->tmu.rate);
+	switch (sw->tmu.mode) {
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		tb_port_tmu_set_unidirectional(down, true);
+		tb_switch_tmu_rate_write(tb_switch_parent(sw), rate);
+		break;
+
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		tb_port_tmu_set_unidirectional(down, false);
+		tb_switch_tmu_rate_write(sw, rate);
+		break;
 
-	tb_switch_set_tmu_mode_params(sw, sw->tmu.rate);
-	tb_port_tmu_set_unidirectional(up, sw->tmu.unidirectional);
+	default:
+		break;
+	}
+
+	tb_switch_set_tmu_mode_params(sw, sw->tmu.mode);
+
+	switch (sw->tmu.mode) {
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		tb_port_tmu_set_unidirectional(up, true);
+		break;
+
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		tb_port_tmu_set_unidirectional(up, false);
+		break;
+
+	default:
+		break;
+	}
 }
 
 static int tb_switch_tmu_change_mode(struct tb_switch *sw)
 {
+	unsigned int rate = tmu_rates[sw->tmu.mode_request];
 	struct tb_port *up, *down;
 	int ret;
 
 	up = tb_upstream_port(sw);
 	down = tb_switch_downstream_port(sw);
-	ret = tb_port_tmu_set_unidirectional(down, sw->tmu.unidirectional_request);
-	if (ret)
-		goto out;
 
-	if (sw->tmu.unidirectional_request)
-		ret = tb_switch_tmu_rate_write(tb_switch_parent(sw),
-					       sw->tmu.rate_request);
-	else
-		ret = tb_switch_tmu_rate_write(sw, sw->tmu.rate_request);
-	if (ret)
-		return ret;
+	/* Program the upstream router downstream facing lane adapter */
+	switch (sw->tmu.mode_request) {
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		ret = tb_port_tmu_set_unidirectional(down, true);
+		if (ret)
+			goto out;
+		ret = tb_switch_tmu_rate_write(tb_switch_parent(sw), rate);
+		if (ret)
+			goto out;
+		break;
+
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		ret = tb_port_tmu_set_unidirectional(down, false);
+		if (ret)
+			goto out;
+		ret = tb_switch_tmu_rate_write(sw, rate);
+		if (ret)
+			goto out;
+		break;
 
-	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.rate_request);
+	default:
+		/* Not allowed to change modes from other than above */
+		return -EINVAL;
+	}
+
+	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.mode_request);
 	if (ret)
 		return ret;
 
-	ret = tb_port_tmu_set_unidirectional(up, sw->tmu.unidirectional_request);
-	if (ret)
-		goto out;
+	/* Program the new mode and the downstream router lane adapter */
+	switch (sw->tmu.mode_request) {
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		ret = tb_port_tmu_set_unidirectional(up, true);
+		if (ret)
+			goto out;
+		break;
+
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		ret = tb_port_tmu_set_unidirectional(up, false);
+		if (ret)
+			goto out;
+		break;
+
+	default:
+		/* Not allowed to change modes from other than above */
+		return -EINVAL;
+	}
 
 	ret = tb_port_tmu_time_sync_enable(down);
 	if (ret)
@@ -637,13 +939,14 @@ out:
  */
 int tb_switch_tmu_enable(struct tb_switch *sw)
 {
-	bool unidirectional = sw->tmu.unidirectional_request;
 	int ret;
 
 	if (tb_switch_tmu_is_enabled(sw))
 		return 0;
 
-	if (tb_switch_is_titan_ridge(sw) && unidirectional) {
+	if (tb_switch_is_titan_ridge(sw) &&
+	    (sw->tmu.mode_request == TB_SWITCH_TMU_MODE_LOWRES ||
+	     sw->tmu.mode_request == TB_SWITCH_TMU_MODE_HIFI_UNI)) {
 		ret = tb_switch_tmu_disable_objections(sw);
 		if (ret)
 			return ret;
@@ -659,19 +962,30 @@ int tb_switch_tmu_enable(struct tb_switch *sw)
 		 * HiFi-Uni/HiFi-BiDir/Normal-Uni or from Normal-Uni to
 		 * HiFi-Uni.
 		 */
-		if (sw->tmu.rate == TB_SWITCH_TMU_RATE_OFF) {
-			if (unidirectional)
+		if (sw->tmu.mode == TB_SWITCH_TMU_MODE_OFF) {
+			switch (sw->tmu.mode_request) {
+			case TB_SWITCH_TMU_MODE_LOWRES:
+			case TB_SWITCH_TMU_MODE_HIFI_UNI:
 				ret = tb_switch_tmu_enable_unidirectional(sw);
-			else
+				break;
+
+			case TB_SWITCH_TMU_MODE_HIFI_BI:
 				ret = tb_switch_tmu_enable_bidirectional(sw);
-			if (ret)
-				return ret;
-		} else if (sw->tmu.rate == TB_SWITCH_TMU_RATE_NORMAL) {
+				break;
+			case TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI:
+				ret = tb_switch_tmu_enable_enhanced(sw);
+				break;
+			default:
+				ret = -EINVAL;
+				break;
+			}
+		} else if (sw->tmu.mode == TB_SWITCH_TMU_MODE_LOWRES ||
+			   sw->tmu.mode == TB_SWITCH_TMU_MODE_HIFI_UNI ||
+			   sw->tmu.mode == TB_SWITCH_TMU_MODE_HIFI_BI) {
 			ret = tb_switch_tmu_change_mode(sw);
-			if (ret)
-				return ret;
+		} else {
+			ret = -EINVAL;
 		}
-		sw->tmu.unidirectional = unidirectional;
 	} else {
 		/*
 		 * Host router port configurations are written as
@@ -679,35 +993,68 @@ int tb_switch_tmu_enable(struct tb_switch *sw)
 		 * of the child node - see above.
 		 * Here only the host router' rate configuration is written.
 		 */
-		ret = tb_switch_tmu_rate_write(sw, sw->tmu.rate_request);
-		if (ret)
-			return ret;
+		ret = tb_switch_tmu_rate_write(sw, tmu_rates[sw->tmu.mode_request]);
 	}
 
-	sw->tmu.rate = sw->tmu.rate_request;
+	if (ret) {
+		tb_sw_warn(sw, "TMU: failed to enable mode %s: %d\n",
+			   tmu_mode_name(sw->tmu.mode_request), ret);
+	} else {
+		sw->tmu.mode = sw->tmu.mode_request;
+		tb_sw_dbg(sw, "TMU: mode set to: %s\n", tmu_mode_name(sw->tmu.mode));
+	}
 
-	tb_sw_dbg(sw, "TMU: mode set to: %s\n", tb_switch_tmu_mode_name(sw));
 	return tb_switch_tmu_set_time_disruption(sw, false);
 }
 
 /**
- * tb_switch_tmu_configure() - Configure the TMU rate and directionality
+ * tb_switch_tmu_configure() - Configure the TMU mode
  * @sw: Router whose mode to change
- * @rate: Rate to configure Off/Normal/HiFi
- * @unidirectional: If uni-directional (bi-directional otherwise)
+ * @mode: Mode to configure
  *
- * Selects the rate of the TMU and directionality (uni-directional or
- * bi-directional). Must be called before tb_switch_tmu_enable().
+ * Selects the TMU mode that is enabled when tb_switch_tmu_enable() is
+ * next called.
  *
- * Returns %0 in success and negative errno otherwise.
+ * Returns %0 in success and negative errno otherwise. Specifically
+ * returns %-EOPNOTSUPP if the requested mode is not possible (not
+ * supported by the router and/or topology).
  */
-int tb_switch_tmu_configure(struct tb_switch *sw, enum tb_switch_tmu_rate rate,
-			    bool unidirectional)
+int tb_switch_tmu_configure(struct tb_switch *sw, enum tb_switch_tmu_mode mode)
 {
-	if (unidirectional && !sw->tmu.has_ucap)
+	switch (mode) {
+	case TB_SWITCH_TMU_MODE_OFF:
+		break;
+
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		if (!sw->tmu.has_ucap)
+			return -EOPNOTSUPP;
+		break;
+
+	case TB_SWITCH_TMU_MODE_HIFI_BI:
+		break;
+
+	case TB_SWITCH_TMU_MODE_MEDRES_ENHANCED_UNI: {
+		const struct tb_switch *parent_sw = tb_switch_parent(sw);
+
+		if (!parent_sw || !tb_switch_tmu_enhanced_is_supported(parent_sw))
+			return -EOPNOTSUPP;
+		if (!tb_switch_tmu_enhanced_is_supported(sw))
+			return -EOPNOTSUPP;
+
+		break;
+	}
+
+	default:
+		tb_sw_warn(sw, "TMU: unsupported mode %u\n", mode);
 		return -EINVAL;
+	}
+
+	if (sw->tmu.mode_request != mode) {
+		tb_sw_dbg(sw, "TMU: mode change %s -> %s requested\n",
+			  tmu_mode_name(sw->tmu.mode), tmu_mode_name(mode));
+		sw->tmu.mode_request = mode;
+	}
 
-	sw->tmu.unidirectional_request = unidirectional;
-	sw->tmu.rate_request = rate;
 	return 0;
 }