#include <linux/of.h>
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
+#include <linux/iopoll.h>
 
 #define CDNS_UART_TTY_NAME     "ttyPS"
 #define CDNS_UART_NAME         "xuartps"
 #define CDNS_UART_MAJOR                0       /* use dynamic node allocation */
 #define CDNS_UART_FIFO_SIZE    64      /* FIFO size */
 #define CDNS_UART_REGISTER_SPACE       0x1000
+#define TX_TIMEOUT             500000
 
 /* Rx Trigger level */
 static int rx_trigger_level = 56;
        unsigned int cval = 0;
        unsigned int baud, minbaud, maxbaud;
        unsigned long flags;
-       unsigned int ctrl_reg, mode_reg;
+       unsigned int ctrl_reg, mode_reg, val;
+       int err;
 
        /* Wait for the transmit FIFO to empty before making changes */
        if (!(readl(port->membase + CDNS_UART_CR) &
                                CDNS_UART_CR_TX_DIS)) {
-               while (!(readl(port->membase + CDNS_UART_SR) &
-                               CDNS_UART_SR_TXEMPTY)) {
-                       cpu_relax();
+               err = readl_poll_timeout(port->membase + CDNS_UART_SR,
+                                        val, (val & CDNS_UART_SR_TXEMPTY),
+                                        1000, TX_TIMEOUT);
+               if (err) {
+                       dev_err(port->dev, "timed out waiting for tx empty");
+                       return;
                }
        }
        spin_lock_irqsave(&port->lock, flags);