if (host->ref_clk_enabled == on)
                return 0;
 
+       ufs_mtk_ref_clk_notify(on, PRE_CHANGE, res);
+
        if (on) {
-               ufs_mtk_ref_clk_notify(on, res);
                ufshcd_writel(hba, REFCLK_REQUEST, REG_UFS_REFCLK_CTRL);
        } else {
                ufshcd_delay_us(host->ref_clk_gating_wait_us, 10);
 
        dev_err(hba->dev, "missing ack of refclk req, reg: 0x%x\n", value);
 
-       ufs_mtk_ref_clk_notify(host->ref_clk_enabled, res);
+       ufs_mtk_ref_clk_notify(host->ref_clk_enabled, POST_CHANGE, res);
 
        return -ETIMEDOUT;
 
        host->ref_clk_enabled = on;
        if (on)
                ufshcd_delay_us(host->ref_clk_ungating_wait_us, 10);
-       else
-               ufs_mtk_ref_clk_notify(on, res);
+
+       ufs_mtk_ref_clk_notify(on, POST_CHANGE, res);
 
        return 0;
 }
 
 #define ufs_mtk_crypto_ctrl(res, enable) \
        ufs_mtk_smc(UFS_MTK_SIP_CRYPTO_CTRL, &(res), enable)
 
-#define ufs_mtk_ref_clk_notify(on, res) \
-       ufs_mtk_smc(UFS_MTK_SIP_REF_CLK_NOTIFICATION, &(res), on)
+#define ufs_mtk_ref_clk_notify(on, stage, res) \
+       ufs_mtk_smc(UFS_MTK_SIP_REF_CLK_NOTIFICATION, &(res), on, stage)
 
 #define ufs_mtk_device_reset_ctrl(high, res) \
        ufs_mtk_smc(UFS_MTK_SIP_DEVICE_RESET, &(res), high)