Signed-off-by: Gil Fine <redacted>
---
 drivers/thunderbolt/tb.c      |   9 +-
 drivers/thunderbolt/tb.h      |  24 ++--
 drivers/thunderbolt/tb_regs.h |   2 +
 drivers/thunderbolt/tmu.c     | 244 ++++++++++++++++++++++++++++------
 4 files changed, 231 insertions(+), 48 deletions(-)

diff --git a/drivers/thunderbolt/tb.c b/drivers/thunderbolt/tb.c
index a231191b06c6..533fe48e85be 100644
--- a/drivers/thunderbolt/tb.c
+++ b/drivers/thunderbolt/tb.c

@@ -221,7 +221,7 @@ static int tb_enable_tmu(struct tb_switch *sw)
 	int ret;
 
 	/* If it is already enabled in correct mode, don't touch it */
-	if (tb_switch_tmu_is_enabled(sw))
+	if (tb_switch_tmu_hifi_is_enabled(sw, sw->tmu.unidirect_request))
 		return 0;
 
 	ret = tb_switch_tmu_disable(sw);

@@ -669,6 +669,7 @@ static void tb_scan_port(struct tb_port *port)
 	tb_switch_lane_bonding_enable(sw);
 	/* Set the link configured */
 	tb_switch_configure_link(sw);
+	tb_switch_tmu_configure(sw, TB_SWITCH_TMU_RATE_HIFI, false);
 
 	if (tb_enable_tmu(sw))
 		tb_sw_warn(sw, "failed to enable TMU\n");

@@ -1375,6 +1376,7 @@ static int tb_start(struct tb *tb)
 		return ret;
 	}
 
+	tb_switch_tmu_configure(tb->root_switch, TB_SWITCH_TMU_RATE_HIFI, false);
 	/* Enable TMU if it is off */
 	tb_switch_tmu_enable(tb->root_switch);
 	/* Full scan to discover devices added before the driver was loaded. */

@@ -1418,6 +1420,11 @@ static void tb_restore_children(struct tb_switch *sw)
 	if (sw->is_unplugged)
 		return;
 
+	/*
+	 * tb_switch_tmu_configure() was already called when the switch was
+	 * added before entering SX/RPM, so no need to call it again

+	 * before enabling TMU.
+	 */
 	if (tb_enable_tmu(sw))
 		tb_sw_warn(sw, "failed to restore TMU configuration\n");
 

diff --git a/drivers/thunderbolt/tb.h b/drivers/thunderbolt/tb.h
index 3fae40670b72..0205361ff89a 100644
--- a/drivers/thunderbolt/tb.h
+++ b/drivers/thunderbolt/tb.h

@@ -89,15 +89,24 @@ enum tb_switch_tmu_rate {
  * @cap: Offset to the TMU capability (%0 if not found)
  * @has_ucap: Does the switch support uni-directional mode
  * @rate: TMU refresh rate related to upstream switch. In case of root
- *	  switch this holds the domain rate.
+ *	  switch this holds the domain rate. Reflects the HW setting.
  * @unidirectional: Is the TMU in uni-directional or bi-directional mode
- *		    related to upstream switch. Don't case for root switch.
+ *		    related to upstream switch. Don't care for root switch.
+ *		    Reflects the HW setting.
+ * @rate_request: TMU new refresh rate related to upstream switch that is
+ *		  requested to be set. In case of root switch, this holds
+ *		  the new domain rate that is requested to be set.
+ * @unidirect_request: Is the new TMU mode: uni-directional or bi-directional
+ *		       that is requested to be set. Related to upstream switch.
+ *		       Don't care for root switch.
  */
 struct tb_switch_tmu {
 	int cap;
 	bool has_ucap;
 	enum tb_switch_tmu_rate rate;
 	bool unidirectional;
+	bool unidirect_request;
+	enum tb_switch_tmu_rate rate_request;
 };
 
 /**

@@ -891,13 +900,10 @@ int tb_switch_tmu_init(struct tb_switch *sw);
 int tb_switch_tmu_post_time(struct tb_switch *sw);
 int tb_switch_tmu_disable(struct tb_switch *sw);
 int tb_switch_tmu_enable(struct tb_switch *sw);
-
-static inline bool tb_switch_tmu_is_enabled(const struct tb_switch *sw)
-{
-	return sw->tmu.rate == TB_SWITCH_TMU_RATE_HIFI &&
-	       !sw->tmu.unidirectional;
-}
-
+bool tb_switch_tmu_hifi_is_enabled(struct tb_switch *sw, bool unidirect);
+void tb_switch_tmu_configure(struct tb_switch *sw,
+			     enum tb_switch_tmu_rate rate,
+			     bool unidirectional);
 int tb_wait_for_port(struct tb_port *port, bool wait_if_unplugged);
 int tb_port_add_nfc_credits(struct tb_port *port, int credits);
 int tb_port_clear_counter(struct tb_port *port, int counter);

diff --git a/drivers/thunderbolt/tb_regs.h b/drivers/thunderbolt/tb_regs.h
index d8ab6c820451..38bc15680d06 100644
--- a/drivers/thunderbolt/tb_regs.h
+++ b/drivers/thunderbolt/tb_regs.h

@@ -305,6 +305,8 @@ struct tb_regs_port_header {
 /* TMU adapter registers */
 #define TMU_ADP_CS_3				0x03
 #define TMU_ADP_CS_3_UDM			BIT(29)
+#define TMU_ADP_CS_6				0x06
+#define TMU_ADP_CS_6_DTS			BIT(1)
 
 /* Lane adapter registers */
 #define LANE_ADP_CS_0				0x00

diff --git a/drivers/thunderbolt/tmu.c b/drivers/thunderbolt/tmu.c
index 039c42a06000..1374813bfc5c 100644
--- a/drivers/thunderbolt/tmu.c
+++ b/drivers/thunderbolt/tmu.c

@@ -115,6 +115,11 @@ static inline int tb_port_tmu_unidirectional_disable(struct tb_port *port)
 	return tb_port_tmu_set_unidirectional(port, false);
 }
 
+static inline int tb_port_tmu_unidirectional_enable(struct tb_port *port)
+{
+	return tb_port_tmu_set_unidirectional(port, true);
+}
+
 static bool tb_port_tmu_is_unidirectional(struct tb_port *port)
 {
 	int ret;

@@ -128,6 +133,23 @@ static bool tb_port_tmu_is_unidirectional(struct tb_port *port)
 	return val & TMU_ADP_CS_3_UDM;
 }
 
+static int tb_port_tmu_time_sync(struct tb_port *port, bool time_sync)
+{
+	u32 val = time_sync ? TMU_ADP_CS_6_DTS : 0;
+
+	return tb_port_tmu_write(port, TMU_ADP_CS_6, TMU_ADP_CS_6_DTS, val);
+}
+
+static int tb_port_tmu_time_sync_disable(struct tb_port *port)
+{
+	return tb_port_tmu_time_sync(port, true);
+}
+
+static int tb_port_tmu_time_sync_enable(struct tb_port *port)
+{
+	return tb_port_tmu_time_sync(port, false);
+}
+
 static int tb_switch_tmu_set_time_disruption(struct tb_switch *sw, bool set)
 {
 	int ret;

@@ -297,6 +319,9 @@ int tb_switch_tmu_post_time(struct tb_switch *sw)
  */
 int tb_switch_tmu_disable(struct tb_switch *sw)
 {
+	struct tb_switch *sw_conf, *parent = tb_switch_parent(sw);

+	bool unidirect = tb_switch_tmu_hifi_is_enabled(sw, true);
+	struct tb_port *up, *down;
 	int ret;
 
 	if (!tb_switch_is_usb4(sw))

@@ -306,21 +331,28 @@ int tb_switch_tmu_disable(struct tb_switch *sw)
 	if (sw->tmu.rate == TB_SWITCH_TMU_RATE_OFF)
 		return 0;
 
-	if (sw->tmu.unidirectional) {
-		struct tb_switch *parent = tb_switch_parent(sw);
-		struct tb_port *up, *down;
+	up = tb_upstream_port(sw);
+	down = tb_port_at(tb_route(sw), parent);
 
-		up = tb_upstream_port(sw);
-		down = tb_port_at(tb_route(sw), parent);
+	if (tb_route(sw)) {
+		sw_conf = unidirect ? parent : sw;
+		tb_switch_tmu_rate_write(sw_conf, TB_SWITCH_TMU_RATE_OFF);

 
-		/* The switch may be unplugged so ignore any errors */
-		tb_port_tmu_unidirectional_disable(up);
-		ret = tb_port_tmu_unidirectional_disable(down);
+		tb_port_tmu_time_sync_disable(up);
+		ret = tb_port_tmu_time_sync_disable(down);
 		if (ret)
 			return ret;
-	}
 
-	tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_OFF);
+		if (unidirect) {
+			/* 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;
+		}
+	} else {
+		tb_switch_tmu_rate_write(sw, TB_SWITCH_TMU_RATE_OFF);
+	}
 
 	sw->tmu.unidirectional = false;
 	sw->tmu.rate = TB_SWITCH_TMU_RATE_OFF;

@@ -329,55 +361,191 @@ int tb_switch_tmu_disable(struct tb_switch *sw)
 	return 0;
 }
 
-/**
- * tb_switch_tmu_enable() - Enable TMU on a switch
- * @sw: Switch whose TMU to enable
- *
- * Enables TMU of a switch to be in bi-directional, HiFi mode. In this mode
- * all tunneling should work.
+/*
+ * This function is called when the previous TMU mode was
+ * TB_SWITCH_TMU_RATE_OFF
  */
-int tb_switch_tmu_enable(struct tb_switch *sw)
+static int __tb_switch_tmu_enable_bidir(struct tb_switch *sw)
 {
+	struct tb_switch *parent = tb_switch_parent(sw);
+	struct tb_port *up, *down;
 	int ret;
 
-	if (!tb_switch_is_usb4(sw))
-		return 0;
+	up = tb_upstream_port(sw);
+	down = tb_port_at(tb_route(sw), parent);
 
-	if (tb_switch_tmu_is_enabled(sw))
-		return 0;
+	ret = tb_port_tmu_unidirectional_enable(up);
+	if (ret)
+		return ret;
 
-	ret = tb_switch_tmu_set_time_disruption(sw, true);
+	ret = tb_port_tmu_unidirectional_enable(down);
+	if (ret)
+		goto out;
+
+	ret = tb_switch_tmu_rate_write(parent, TB_SWITCH_TMU_RATE_HIFI);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_time_sync_enable(up);
+	if (ret)
+		goto out;
+
+	ret = tb_port_tmu_time_sync_enable(down);
+	if (ret)
+		goto out;
+
+	return 0;
+out:
+	/*
+	 * In case of any failure in one of the steps, get back to the
+	 * TMU configurations in OFF mode. In case of additional failures in
+	 * the functions below, ignore them since we already report a failure.
+	 */
+	tb_port_tmu_time_sync_disable(up);
+	tb_switch_tmu_rate_write(parent, TB_SWITCH_TMU_RATE_OFF);
+	tb_port_tmu_unidirectional_disable(down);
+	tb_port_tmu_unidirectional_disable(up);