#include <linux/of.h>
 #include <linux/of_device.h>
 #include <linux/of_dma.h>
+#include <linux/pinctrl/consumer.h>
 #include <linux/serial_core.h>
 #include <linux/slab.h>
 #include <linux/tty_flip.h>
        sport->lpuart_dma_rx_use = false;
 }
 
+static void lpuart_hw_setup(struct lpuart_port *sport)
+{
+       unsigned long flags;
+
+       spin_lock_irqsave(&sport->port.lock, flags);
+
+       lpuart_setup_watermark_enable(sport);
+
+       lpuart_rx_dma_startup(sport);
+       lpuart_tx_dma_startup(sport);
+
+       spin_unlock_irqrestore(&sport->port.lock, flags);
+}
+
 static int lpuart_startup(struct uart_port *port)
 {
        struct lpuart_port *sport = container_of(port, struct lpuart_port, port);
-       unsigned long flags;
        unsigned char temp;
 
        /* determine FIFO size and enable FIFO mode */
                                            UARTPFIFO_FIFOSIZE_MASK);
 
        lpuart_request_dma(sport);
-
-       spin_lock_irqsave(&sport->port.lock, flags);
-
-       lpuart_setup_watermark_enable(sport);
-
-       lpuart_rx_dma_startup(sport);
-       lpuart_tx_dma_startup(sport);
-
-       spin_unlock_irqrestore(&sport->port.lock, flags);
+       lpuart_hw_setup(sport);
 
        return 0;
 }
        lpuart32_write(&sport->port, temp, UARTCTRL);
 }
 
+static void lpuart32_hw_setup(struct lpuart_port *sport)
+{
+       unsigned long flags;
+
+       spin_lock_irqsave(&sport->port.lock, flags);
+
+       lpuart32_setup_watermark_enable(sport);
+
+       lpuart_rx_dma_startup(sport);
+       lpuart_tx_dma_startup(sport);
+
+       lpuart32_configure(sport);
+
+       spin_unlock_irqrestore(&sport->port.lock, flags);
+}
+
 static int lpuart32_startup(struct uart_port *port)
 {
        struct lpuart_port *sport = container_of(port, struct lpuart_port, port);
-       unsigned long flags;
        unsigned long temp;
 
        /* determine FIFO size */
        }
 
        lpuart_request_dma(sport);
+       lpuart32_hw_setup(sport);
 
-       spin_lock_irqsave(&sport->port.lock, flags);
-
-       lpuart32_setup_watermark_enable(sport);
-
-       lpuart_rx_dma_startup(sport);
-       lpuart_tx_dma_startup(sport);
-
-       lpuart32_configure(sport);
-
-       spin_unlock_irqrestore(&sport->port.lock, flags);
        return 0;
 }
 
        return 0;
 }
 
-static int __maybe_unused lpuart_suspend(struct device *dev)
+static void serial_lpuart_enable_wakeup(struct lpuart_port *sport, bool on)
 {
-       struct lpuart_port *sport = dev_get_drvdata(dev);
-       unsigned long temp;
-       bool irq_wake;
+       unsigned int val, baud;
 
        if (lpuart_is_32(sport)) {
-               /* disable Rx/Tx and interrupts */
-               temp = lpuart32_read(&sport->port, UARTCTRL);
-               temp &= ~(UARTCTRL_TE | UARTCTRL_TIE | UARTCTRL_TCIE);
-               lpuart32_write(&sport->port, temp, UARTCTRL);
+               val = lpuart32_read(&sport->port, UARTCTRL);
+               baud = lpuart32_read(&sport->port, UARTBAUD);
+               if (on) {
+                       /* set rx_watermark to 0 in wakeup source mode */
+                       lpuart32_write(&sport->port, 0, UARTWATER);
+                       val |= UARTCTRL_RIE;
+                       /* clear RXEDGIF flag before enable RXEDGIE interrupt */
+                       lpuart32_write(&sport->port, UARTSTAT_RXEDGIF, UARTSTAT);
+                       baud |= UARTBAUD_RXEDGIE;
+               } else {
+                       val &= ~UARTCTRL_RIE;
+                       baud &= ~UARTBAUD_RXEDGIE;
+               }
+               lpuart32_write(&sport->port, val, UARTCTRL);
+               lpuart32_write(&sport->port, baud, UARTBAUD);
        } else {
-               /* disable Rx/Tx and interrupts */
-               temp = readb(sport->port.membase + UARTCR2);
-               temp &= ~(UARTCR2_TE | UARTCR2_TIE | UARTCR2_TCIE);
-               writeb(temp, sport->port.membase + UARTCR2);
+               val = readb(sport->port.membase + UARTCR2);
+               if (on)
+                       val |= UARTCR2_RIE;
+               else
+                       val &= ~UARTCR2_RIE;
+               writeb(val, sport->port.membase + UARTCR2);
        }
+}
 
-       uart_suspend_port(&lpuart_reg, &sport->port);
+static bool lpuart_uport_is_active(struct lpuart_port *sport)
+{
+       struct tty_port *port = &sport->port.state->port;
+       struct tty_struct *tty;
+       struct device *tty_dev;
+       int may_wake = 0;
 
-       /* uart_suspend_port() might set wakeup flag */
-       irq_wake = irqd_is_wakeup_set(irq_get_irq_data(sport->port.irq));
+       tty = tty_port_tty_get(port);
+       if (tty) {
+               tty_dev = tty->dev;
+               may_wake = device_may_wakeup(tty_dev);
+               tty_kref_put(tty);
+       }
 
-       if (sport->lpuart_dma_rx_use) {
-               /*
-                * EDMA driver during suspend will forcefully release any
-                * non-idle DMA channels. If port wakeup is enabled or if port
-                * is console port or 'no_console_suspend' is set the Rx DMA
-                * cannot resume as expected, hence gracefully release the
-                * Rx DMA path before suspend and start Rx DMA path on resume.
-                */
-               if (irq_wake) {
-                       del_timer_sync(&sport->lpuart_timer);
-                       lpuart_dma_rx_free(&sport->port);
-               }
+       if ((tty_port_initialized(port) && may_wake) ||
+           (!console_suspend_enabled && uart_console(&sport->port)))
+               return true;
+
+       return false;
+}
+
+static int __maybe_unused lpuart_suspend_noirq(struct device *dev)
+{
+       struct lpuart_port *sport = dev_get_drvdata(dev);
+       bool irq_wake = irqd_is_wakeup_set(irq_get_irq_data(sport->port.irq));
+
+       if (lpuart_uport_is_active(sport))
+               serial_lpuart_enable_wakeup(sport, !!irq_wake);
+
+       pinctrl_pm_select_sleep_state(dev);
+
+       return 0;
+}
+
+static int __maybe_unused lpuart_resume_noirq(struct device *dev)
+{
+       struct lpuart_port *sport = dev_get_drvdata(dev);
+       unsigned int val;
+
+       pinctrl_pm_select_default_state(dev);
+
+       if (lpuart_uport_is_active(sport)) {
+               serial_lpuart_enable_wakeup(sport, false);
 
-               /* Disable Rx DMA to use UART port as wakeup source */
+               /* clear the wakeup flags */
                if (lpuart_is_32(sport)) {
-                       temp = lpuart32_read(&sport->port, UARTBAUD);
-                       lpuart32_write(&sport->port, temp & ~UARTBAUD_RDMAE,
-                                      UARTBAUD);
-               } else {
-                       writeb(readb(sport->port.membase + UARTCR5) &
-                              ~UARTCR5_RDMAS, sport->port.membase + UARTCR5);
+                       val = lpuart32_read(&sport->port, UARTSTAT);
+                       lpuart32_write(&sport->port, val, UARTSTAT);
                }
        }
 
-       if (sport->lpuart_dma_tx_use) {
-               sport->dma_tx_in_progress = false;
-               dmaengine_terminate_all(sport->dma_tx_chan);
-       }
-
-       if (sport->port.suspended && !irq_wake)
-               lpuart_disable_clks(sport);
-
        return 0;
 }
 
-static int __maybe_unused lpuart_resume(struct device *dev)
+static int __maybe_unused lpuart_suspend(struct device *dev)
 {
        struct lpuart_port *sport = dev_get_drvdata(dev);
-       bool irq_wake = irqd_is_wakeup_set(irq_get_irq_data(sport->port.irq));
+       unsigned long temp, flags;
 
-       if (sport->port.suspended && !irq_wake)
-               lpuart_enable_clks(sport);
+       uart_suspend_port(&lpuart_reg, &sport->port);
 
-       if (lpuart_is_32(sport))
-               lpuart32_setup_watermark_enable(sport);
-       else
-               lpuart_setup_watermark_enable(sport);
+       if (lpuart_uport_is_active(sport)) {
+               spin_lock_irqsave(&sport->port.lock, flags);
+               if (lpuart_is_32(sport)) {
+                       /* disable Rx/Tx and interrupts */
+                       temp = lpuart32_read(&sport->port, UARTCTRL);
+                       temp &= ~(UARTCTRL_TE | UARTCTRL_TIE | UARTCTRL_TCIE);
+                       lpuart32_write(&sport->port, temp, UARTCTRL);
+               } else {
+                       /* disable Rx/Tx and interrupts */
+                       temp = readb(sport->port.membase + UARTCR2);
+                       temp &= ~(UARTCR2_TE | UARTCR2_TIE | UARTCR2_TCIE);
+                       writeb(temp, sport->port.membase + UARTCR2);
+               }
+               spin_unlock_irqrestore(&sport->port.lock, flags);
 
-       if (sport->lpuart_dma_rx_use) {
-               if (irq_wake) {
-                       if (!lpuart_start_rx_dma(sport))
-                               rx_dma_timer_init(sport);
-                       else
-                               sport->lpuart_dma_rx_use = false;
+               if (sport->lpuart_dma_rx_use) {
+                       /*
+                        * EDMA driver during suspend will forcefully release any
+                        * non-idle DMA channels. If port wakeup is enabled or if port
+                        * is console port or 'no_console_suspend' is set the Rx DMA
+                        * cannot resume as expected, hence gracefully release the
+                        * Rx DMA path before suspend and start Rx DMA path on resume.
+                        */
+                       del_timer_sync(&sport->lpuart_timer);
+                       lpuart_dma_rx_free(&sport->port);
+
+                       /* Disable Rx DMA to use UART port as wakeup source */
+                       spin_lock_irqsave(&sport->port.lock, flags);
+                       if (lpuart_is_32(sport)) {
+                               temp = lpuart32_read(&sport->port, UARTBAUD);
+                               lpuart32_write(&sport->port, temp & ~UARTBAUD_RDMAE,
+                                              UARTBAUD);
+                       } else {
+                               writeb(readb(sport->port.membase + UARTCR5) &
+                                      ~UARTCR5_RDMAS, sport->port.membase + UARTCR5);
+                       }
+                       spin_unlock_irqrestore(&sport->port.lock, flags);
+               }
+
+               if (sport->lpuart_dma_tx_use) {
+                       spin_lock_irqsave(&sport->port.lock, flags);
+                       if (lpuart_is_32(sport)) {
+                               temp = lpuart32_read(&sport->port, UARTBAUD);
+                               temp &= ~UARTBAUD_TDMAE;
+                               lpuart32_write(&sport->port, temp, UARTBAUD);
+                       } else {
+                               temp = readb(sport->port.membase + UARTCR5);
+                               temp &= ~UARTCR5_TDMAS;
+                               writeb(temp, sport->port.membase + UARTCR5);
+                       }
+                       spin_unlock_irqrestore(&sport->port.lock, flags);
+                       sport->dma_tx_in_progress = false;
+                       dmaengine_terminate_all(sport->dma_tx_chan);
                }
        }
 
-       lpuart_tx_dma_startup(sport);
+       return 0;
+}
 
-       if (lpuart_is_32(sport))
-               lpuart32_configure(sport);
+static void lpuart_console_fixup(struct lpuart_port *sport)
+{
+       struct tty_port *port = &sport->port.state->port;
+       struct uart_port *uport = &sport->port;
+       struct ktermios termios;
+
+       /* i.MX7ULP enter VLLS mode that lpuart module power off and registers
+        * all lost no matter the port is wakeup source.
+        * For console port, console baud rate setting lost and print messy
+        * log when enable the console port as wakeup source. To avoid the
+        * issue happen, user should not enable uart port as wakeup source
+        * in VLLS mode, or restore console setting here.
+        */
+       if (is_imx7ulp_lpuart(sport) && lpuart_uport_is_active(sport) &&
+           console_suspend_enabled && uart_console(&sport->port)) {
+
+               mutex_lock(&port->mutex);
+               memset(&termios, 0, sizeof(struct ktermios));
+               termios.c_cflag = uport->cons->cflag;
+               if (port->tty && termios.c_cflag == 0)
+                       termios = port->tty->termios;
+               uport->ops->set_termios(uport, &termios, NULL);
+               mutex_unlock(&port->mutex);
+       }
+}
+
+static int __maybe_unused lpuart_resume(struct device *dev)
+{
+       struct lpuart_port *sport = dev_get_drvdata(dev);
+
+       if (lpuart_uport_is_active(sport)) {
+               if (lpuart_is_32(sport))
+                       lpuart32_hw_setup(sport);
+               else
+                       lpuart_hw_setup(sport);
+       }
 
+       lpuart_console_fixup(sport);
        uart_resume_port(&lpuart_reg, &sport->port);
 
        return 0;
 }
 
-static SIMPLE_DEV_PM_OPS(lpuart_pm_ops, lpuart_suspend, lpuart_resume);
+static const struct dev_pm_ops lpuart_pm_ops = {
+       SET_NOIRQ_SYSTEM_SLEEP_PM_OPS(lpuart_suspend_noirq,
+                                     lpuart_resume_noirq)
+       SET_SYSTEM_SLEEP_PM_OPS(lpuart_suspend, lpuart_resume)
+};
 
 static struct platform_driver lpuart_driver = {
        .probe          = lpuart_probe,