While ufs_mtk_wait_idle state has some code smells for me (the
VS_HCE_BASE early exit seems racey at best), it can still benefit from
some general cleanup to make the code flow less convoluted.

Use the iopoll helpers, for one, and specifically the one that sleeps
and does not busy delay, as it's being done for up to 5ms.

The register read is split out to a helper function that branches
between new and old style flow.

Every called uses the same 5ms timeout value, so there is no point in
making this a parameter. Just assume a 5ms timeout in the function.

Reviewed-by: AngeloGioacchino Del Regno <angelogioacchino.delregno@collabora.com>
Signed-off-by: Nicolas Frattaroli <nicolas.frattaroli@collabora.com>
---
 drivers/ufs/host/ufs-mediatek.c | 71 +++++++++++++++++------------------------
 1 file changed, 30 insertions(+), 41 deletions(-)

diff --git a/drivers/ufs/host/ufs-mediatek.c b/drivers/ufs/host/ufs-mediatek.c
index 2b1f26b55782..b5c75d228c85 100644
--- a/drivers/ufs/host/ufs-mediatek.c
+++ b/drivers/ufs/host/ufs-mediatek.c
@@ -10,6 +10,7 @@
 #include <linux/bitfield.h>
 #include <linux/clk.h>
 #include <linux/delay.h>
+#include <linux/iopoll.h>
 #include <linux/module.h>
 #include <linux/of.h>
 #include <linux/of_address.h>
@@ -380,51 +381,39 @@ static void ufs_mtk_dbg_sel(struct ufs_hba *hba)
 	}
 }
 
-static int ufs_mtk_wait_idle_state(struct ufs_hba *hba,
-			    unsigned long retry_ms)
+static u32 ufs_mtk_read_state(struct ufs_hba *hba, bool old_style)
 {
-	u64 timeout, time_checked;
-	u32 val, sm;
-	bool wait_idle;
-	struct ufs_mtk_host *host = ufshcd_get_variant(hba);
-
-	/* cannot use plain ktime_get() in suspend */
-	timeout = ktime_get_mono_fast_ns() + retry_ms * 1000000UL;
-
-	/* wait a specific time after check base */
-	udelay(10);
-	wait_idle = false;
+	u32 val;
 
-	do {
-		time_checked = ktime_get_mono_fast_ns();
-		if (host->legacy_ip_ver || host->ip_ver < IP_VER_MT6899) {
-			ufs_mtk_dbg_sel(hba);
-			val = ufshcd_readl(hba, REG_UFS_PROBE);
-		} else {
-			val = ufshcd_readl(hba, REG_UFS_UFS_MMIO_OTSD_CTRL);
-			val = val >> 16;
-		}
+	if (old_style) {
+		ufs_mtk_dbg_sel(hba);
+		val = ufshcd_readl(hba, REG_UFS_PROBE);
+	} else {
+		val = ufshcd_readl(hba, REG_UFS_UFS_MMIO_OTSD_CTRL) >> 16;
+	}
 
-		sm = val & 0x1f;
+	return FIELD_GET(0x1f, val);
+}
 
-		/*
-		 * if state is in H8 enter and H8 enter confirm
-		 * wait until return to idle state.
-		 */
-		if ((sm >= VS_HIB_ENTER) && (sm <= VS_HIB_EXIT)) {
-			wait_idle = true;
-			udelay(50);
-			continue;
-		} else if (!wait_idle)
-			break;
+static int ufs_mtk_wait_idle_state(struct ufs_hba *hba)
+{
+	struct ufs_mtk_host *host = ufshcd_get_variant(hba);
+	bool old_style = (host->legacy_ip_ver || host->ip_ver < IP_VER_MT6899);
+	u32 val;
+	int ret;
 
-		if (wait_idle && (sm == VS_HCE_BASE))
-			break;
-	} while (time_checked < timeout);
+	/* If the device is already in the base state after 10us, don't wait. */
+	udelay(10);
+	if (ufs_mtk_read_state(hba, old_style) == VS_HCE_BASE)
+		return 0;
 
-	if (wait_idle && sm != VS_HCE_BASE) {
-		dev_info(hba->dev, "wait idle tmo: 0x%x\n", val);
-		return -ETIMEDOUT;
+	/* Poll to wait for idle */
+	ret = read_poll_timeout(ufs_mtk_read_state, val,
+				(val < VS_HIB_ENTER || val > VS_HIB_EXIT), 50,
+				5 * USEC_PER_MSEC, false, hba, old_style);
+	if (ret) {
+		dev_err(hba->dev, "Timed out waiting for idle state, val = 0x%x\n", val);
+		return ret;
 	}
 
 	return 0;
@@ -1389,7 +1378,7 @@ static int ufs_mtk_auto_hibern8_disable(struct ufs_hba *hba)
 	ufshcd_writel(hba, 0, REG_AUTO_HIBERNATE_IDLE_TIMER);
 
 	/* wait host return to idle state when auto-hibern8 off */
-	ret = ufs_mtk_wait_idle_state(hba, 5);
+	ret = ufs_mtk_wait_idle_state(hba);
 	if (ret)
 		goto out;
 
@@ -1593,7 +1582,7 @@ static int ufs_mtk_link_set_hpm(struct ufs_hba *hba)
 		return err;
 
 	/* Check link state to make sure exit h8 success */
-	err = ufs_mtk_wait_idle_state(hba, 5);
+	err = ufs_mtk_wait_idle_state(hba);
 	if (err) {
 		dev_err(hba->dev, "Failed to wait for idle: %pe\n", ERR_PTR(err));
 		return err;

-- 
2.53.0

On Mon, 2026-02-16 at 14:37 +0100, Nicolas Frattaroli wrote:
> While ufs_mtk_wait_idle state has some code smells for me (the
> VS_HCE_BASE early exit seems racey at best), it can still benefit
> from
> some general cleanup to make the code flow less convoluted.
> 
> Use the iopoll helpers, for one, and specifically the one that sleeps
> and does not busy delay, as it's being done for up to 5ms.
> 
> The register read is split out to a helper function that branches
> between new and old style flow.
> 
> Every called uses the same 5ms timeout value, so there is no point in
> making this a parameter. Just assume a 5ms timeout in the function.
> 
> Reviewed-by: AngeloGioacchino Del Regno
> <angelogioacchino.delregno@collabora.com>
> Signed-off-by: Nicolas Frattaroli <nicolas.frattaroli@collabora.com>

Hi Nicolas,

The logic is quite different.

Previously, the check was:
If (sm >= VS_HIB_ENTER) and (sm <= VS_HIB_EXIT), then wait for sm to
reach VS_HCE_BASE.
If not ((val < VS_HIB_ENTER) and (val > VS_HIB_EXIT)), then exit the
loop directly.

Now, the logic is:
If sm == VS_HCE_BASE, return 0.
If (sm >= VS_HIB_ENTER) and (sm <= VS_HIB_EXIT), then wait for sm to 
transition to any other state and exit the poll.

Thanks
Peter

Il 25/02/26 11:32, Peter Wang (王信友) ha scritto:
> On Mon, 2026-02-16 at 14:37 +0100, Nicolas Frattaroli wrote:
>> While ufs_mtk_wait_idle state has some code smells for me (the
>> VS_HCE_BASE early exit seems racey at best), it can still benefit
>> from
>> some general cleanup to make the code flow less convoluted.
>>
>> Use the iopoll helpers, for one, and specifically the one that sleeps
>> and does not busy delay, as it's being done for up to 5ms.
>>
>> The register read is split out to a helper function that branches
>> between new and old style flow.
>>
>> Every called uses the same 5ms timeout value, so there is no point in
>> making this a parameter. Just assume a 5ms timeout in the function.
>>
>> Reviewed-by: AngeloGioacchino Del Regno
>> <angelogioacchino.delregno@collabora.com>
>> Signed-off-by: Nicolas Frattaroli <nicolas.frattaroli@collabora.com>
> 
> Hi Nicolas,
> 
> The logic is quite different.
> 
> Previously, the check was:
> If (sm >= VS_HIB_ENTER) and (sm <= VS_HIB_EXIT), then wait for sm to
> reach VS_HCE_BASE.
> If not ((val < VS_HIB_ENTER) and (val > VS_HIB_EXIT)), then exit the
> loop directly.
> 
> Now, the logic is:
> If sm == VS_HCE_BASE, return 0.
> If (sm >= VS_HIB_ENTER) and (sm <= VS_HIB_EXIT), then wait for sm to
> transition to any other state and exit the poll.
> 

The logic is practically the same, except it's proper now.

There's no need to wait for `sm = VS_HCE_BASE` if `sm == VS_HCE_BASE`.

In other words, just read the comment that Nicolas has put in the code:
/* If the device is already in the base state after 10us, don't wait. */

...because there's no need to wait for the device to get to BASE state,
if the device is *already* in BASE state. It's pretty obvious stuff here.

Regards,
Angelo

> Thanks
> Peter

On Wed, 2026-02-25 at 13:33 +0100, AngeloGioacchino Del Regno wrote:
> The logic is practically the same, except it's proper now.
> 
> There's no need to wait for `sm = VS_HCE_BASE` if `sm ==
> VS_HCE_BASE`.
> 
> In other words, just read the comment that Nicolas has put in the
> code:
> /* If the device is already in the base state after 10us, don't wait.
> */
> 
> ...because there's no need to wait for the device to get to BASE
> state,
> if the device is *already* in BASE state. It's pretty obvious stuff
> here.
> 
> Regards,
> Angelo

Hi AngeloGioacchino,

You missed my point.
Yes, it's true that if the device is already in the base
state after 10 μs, you don't need to wait.
But if the device is in another state after 10 μs, such 
as VS_HCE_DME_RESET, you also don't need to wait.

Besides, if after 10 μs the read state is between 
VS_HIB_EXIT and VS_HIB_ENTER, then you should wait until 
it reaches VS_HCE_BASE, not another state like VS_HCE_DME_RESET.

Thanks.
Peter