1 files changed, 500 insertions, 213 deletions

diff --git a/drivers/thunderbolt/tmu.c b/drivers/thunderbolt/tmu.c
index 626aca3124b1..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;
 }
 
@@ -308,7 +466,7 @@ int tb_switch_tmu_post_time(struct tb_switch *sw)
 		return ret;
 
 	for (i = 0; i < ARRAY_SIZE(gm_local_time); i++)
-		tb_sw_dbg(root_switch, "local_time[%d]=0x%08x\n", i,
+		tb_sw_dbg(root_switch, "TMU: local_time[%d]=0x%08x\n", i,
 			  gm_local_time[i]);
 
 	/* Convert to nanoseconds (drop fractional part) */
@@ -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
@@ -383,26 +558,15 @@ out:
  */
 int tb_switch_tmu_disable(struct tb_switch *sw)
 {
-	/*
-	 * No need to disable TMU on devices that don't support CLx since
-	 * on these devices e.g. Alpine Ridge and earlier, the TMU mode
-	 * HiFi bi-directional is enabled by default and we don't change it.
-	 */
-	if (!tb_switch_is_clx_supported(sw))
-		return 0;
-
 	/* 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_switch *parent = tb_switch_parent(sw);
 		struct tb_port *down, *up;
 		int ret;
 
-		down = tb_port_at(tb_route(sw), parent);
+		down = tb_switch_downstream_port(sw);
 		up = tb_upstream_port(sw);
 		/*
 		 * In case of uni-directional time sync, TMU handshake is
@@ -415,37 +579,49 @@ 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)
 {
-	struct tb_switch *parent = tb_switch_parent(sw);
+	unsigned int rate = tmu_rates[TB_SWITCH_TMU_MODE_OFF];
 	struct tb_port *down, *up;
 
-	down = tb_port_at(tb_route(sw), parent);
+	down = tb_switch_downstream_port(sw);
 	up = tb_upstream_port(sw);
 	/*
 	 * In case of any failure in one of the steps when setting
@@ -456,28 +632,38 @@ 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(parent, 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)
+static int tb_switch_tmu_enable_bidirectional(struct tb_switch *sw)
 {
-	struct tb_switch *parent = tb_switch_parent(sw);
 	struct tb_port *up, *down;
 	int ret;
 
 	up = tb_upstream_port(sw);
-	down = tb_port_at(tb_route(sw), parent);
+	down = tb_switch_downstream_port(sw);
 
 	ret = tb_port_tmu_unidirectional_disable(up);
 	if (ret)
@@ -487,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;
 
@@ -502,12 +688,14 @@ 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;
 }
 
-static int tb_switch_tmu_objection_mask(struct tb_switch *sw)
+/* Only needed for Titan Ridge */
+static int tb_switch_tmu_disable_objections(struct tb_switch *sw)
 {
+	struct tb_port *up = tb_upstream_port(sw);
 	u32 val;
 	int ret;
 
@@ -518,36 +706,34 @@ static int tb_switch_tmu_objection_mask(struct tb_switch *sw)
 
 	val &= ~TB_TIME_VSEC_3_CS_9_TMU_OBJ_MASK;
 
-	return tb_sw_write(sw, &val, TB_CFG_SWITCH,
-			   sw->cap_vsec_tmu + TB_TIME_VSEC_3_CS_9, 1);
-}
-
-static int tb_switch_tmu_unidirectional_enable(struct tb_switch *sw)
-{
-	struct tb_port *up = tb_upstream_port(sw);
+	ret = tb_sw_write(sw, &val, TB_CFG_SWITCH,
+			  sw->cap_vsec_tmu + TB_TIME_VSEC_3_CS_9, 1);
+	if (ret)
+		return ret;
 
 	return tb_port_tmu_write(up, TMU_ADP_CS_6,
 				 TMU_ADP_CS_6_DISABLE_TMU_OBJ_MASK,
-				 TMU_ADP_CS_6_DISABLE_TMU_OBJ_MASK);
+				 TMU_ADP_CS_6_DISABLE_TMU_OBJ_CL1 |
+				 TMU_ADP_CS_6_DISABLE_TMU_OBJ_CL2);
 }
 
 /*
  * 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)
+static int tb_switch_tmu_enable_unidirectional(struct tb_switch *sw)
 {
-	struct tb_switch *parent = tb_switch_parent(sw);
 	struct tb_port *up, *down;
 	int ret;
 
 	up = tb_upstream_port(sw);
-	down = tb_port_at(tb_route(sw), parent);
-	ret = tb_switch_tmu_rate_write(parent, sw->tmu.rate_request);
+	down = tb_switch_downstream_port(sw);
+	ret = tb_switch_tmu_rate_write(tb_switch_parent(sw),
+				       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;
 
@@ -570,16 +756,65 @@ 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)
+static void tb_switch_tmu_change_mode_prev(struct tb_switch *sw)
 {
-	struct tb_switch *parent = tb_switch_parent(sw);
+	unsigned int rate = tmu_rates[sw->tmu.mode];
 	struct tb_port *down, *up;
 
-	down = tb_port_at(tb_route(sw), parent);
+	down = tb_switch_downstream_port(sw);
 	up = tb_upstream_port(sw);
 	/*
 	 * In case of any failure in one of the steps when change mode,
@@ -587,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(parent, 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;
+
+	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;
 
-	tb_switch_set_tmu_mode_params(sw, sw->tmu.rate);
-	tb_port_tmu_set_unidirectional(up, sw->tmu.unidirectional);
+	default:
+		break;
+	}
 }
 
-static int __tb_switch_tmu_change_mode(struct tb_switch *sw)
+static int tb_switch_tmu_change_mode(struct tb_switch *sw)
 {
-	struct tb_switch *parent = tb_switch_parent(sw);
+	unsigned int rate = tmu_rates[sw->tmu.mode_request];
 	struct tb_port *up, *down;
 	int ret;
 
 	up = tb_upstream_port(sw);
-	down = tb_port_at(tb_route(sw), parent);
-	ret = tb_port_tmu_set_unidirectional(down, sw->tmu.unidirectional_request);
-	if (ret)
-		goto out;
+	down = tb_switch_downstream_port(sw);
 
-	if (sw->tmu.unidirectional_request)
-		ret = tb_switch_tmu_rate_write(parent, 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;
 
-	ret = tb_switch_set_tmu_mode_params(sw, sw->tmu.rate_request);
+	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;
+
+	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)
@@ -635,7 +925,7 @@ static int __tb_switch_tmu_change_mode(struct tb_switch *sw)
 	return 0;
 
 out:
-	__tb_switch_tmu_change_mode_prev(sw);
+	tb_switch_tmu_change_mode_prev(sw);
 	return ret;
 }
 
@@ -643,45 +933,21 @@ out:
  * tb_switch_tmu_enable() - Enable TMU on a router
  * @sw: Router whose TMU to enable
  *
- * Enables TMU of a router to be in uni-directional Normal/HiFi
- * or bi-directional HiFi mode. Calling tb_switch_tmu_configure() is required
- * before calling this function, to select the mode Normal/HiFi and
- * directionality (uni-directional/bi-directional).
- * In HiFi mode all tunneling should work. In Normal mode, DP tunneling can't
- * work. Uni-directional mode is required for CLx (Link Low-Power) to work.
+ * Enables TMU of a router to be in uni-directional Normal/HiFi or
+ * bi-directional HiFi mode. Calling tb_switch_tmu_configure() is
+ * required before calling this function.
  */
 int tb_switch_tmu_enable(struct tb_switch *sw)
 {
-	bool unidirectional = sw->tmu.unidirectional_request;
 	int ret;
 
-	if (unidirectional && !sw->tmu.has_ucap)
-		return -EOPNOTSUPP;
-
-	/*
-	 * No need to enable TMU on devices that don't support CLx since on
-	 * these devices e.g. Alpine Ridge and earlier, the TMU mode HiFi
-	 * bi-directional is enabled by default.
-	 */
-	if (!tb_switch_is_clx_supported(sw))
-		return 0;
-
-	if (tb_switch_tmu_is_enabled(sw, sw->tmu.unidirectional_request))
+	if (tb_switch_tmu_is_enabled(sw))
 		return 0;
 
-	if (tb_switch_is_titan_ridge(sw) && unidirectional) {
-		/*
-		 * Titan Ridge supports CL0s and CL1 only. CL0s and CL1 are
-		 * enabled and supported together.
-		 */
-		if (!tb_switch_is_clx_enabled(sw, TB_CL1))
-			return -EOPNOTSUPP;
-
-		ret = tb_switch_tmu_objection_mask(sw);
-		if (ret)
-			return ret;
-
-		ret = tb_switch_tmu_unidirectional_enable(sw);
+	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;
 	}
@@ -696,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)
-				ret = __tb_switch_tmu_enable_unidirectional(sw);
-			else
-				ret = __tb_switch_tmu_enable_bidirectional(sw);
-			if (ret)
-				return ret;
-		} else if (sw->tmu.rate == TB_SWITCH_TMU_RATE_NORMAL) {
-			ret = __tb_switch_tmu_change_mode(sw);
-			if (ret)
-				return ret;
+		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);
+				break;
+
+			case TB_SWITCH_TMU_MODE_HIFI_BI:
+				ret = tb_switch_tmu_enable_bidirectional(sw);
+				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);
+		} else {
+			ret = -EINVAL;
 		}
-		sw->tmu.unidirectional = unidirectional;
 	} else {
 		/*
 		 * Host router port configurations are written as
@@ -716,58 +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. Specifically
+ * returns %-EOPNOTSUPP if the requested mode is not possible (not
+ * supported by the router and/or topology).
  */
-void 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)
 {
-	sw->tmu.unidirectional_request = unidirectional;
-	sw->tmu.rate_request = rate;
-}
+	switch (mode) {
+	case TB_SWITCH_TMU_MODE_OFF:
+		break;
 
-static int tb_switch_tmu_config_enable(struct device *dev, void *rate)
-{
-	if (tb_is_switch(dev)) {
-		struct tb_switch *sw = tb_to_switch(dev);
+	case TB_SWITCH_TMU_MODE_LOWRES:
+	case TB_SWITCH_TMU_MODE_HIFI_UNI:
+		if (!sw->tmu.has_ucap)
+			return -EOPNOTSUPP;
+		break;
 
-		tb_switch_tmu_configure(sw, *(enum tb_switch_tmu_rate *)rate,
-					tb_switch_is_clx_enabled(sw, TB_CL1));
-		if (tb_switch_tmu_enable(sw))
-			tb_sw_dbg(sw, "fail switching TMU mode for 1st depth router\n");
+	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;
 	}
 
-	return 0;
-}
+	default:
+		tb_sw_warn(sw, "TMU: unsupported mode %u\n", mode);
+		return -EINVAL;
+	}
 
-/**
- * tb_switch_enable_tmu_1st_child - Configure and enable TMU for 1st chidren
- * @sw: The router to configure and enable it's children TMU
- * @rate: Rate of the TMU to configure the router's chidren to
- *
- * Configures and enables the TMU mode of 1st depth children of the specified
- * router to the specified rate.
- */
-void tb_switch_enable_tmu_1st_child(struct tb_switch *sw,
-				    enum tb_switch_tmu_rate rate)
-{
-	device_for_each_child(&sw->dev, &rate,
-			      tb_switch_tmu_config_enable);
+	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;
+	}
+
+	return 0;
 }