#include <linux/i2c.h>
 #include <linux/interrupt.h>
 #include <linux/io.h>
+#include <linux/iopoll.h>
 #include <linux/module.h>
 #include <linux/of_device.h>
 
 #define OWL_I2C_FIFOCTL_TFR    BIT(2)
 
 /* I2Cc_FIFOSTAT Bit Mask */
+#define OWL_I2C_FIFOSTAT_CECB  BIT(0)
 #define OWL_I2C_FIFOSTAT_RNB   BIT(1)
 #define OWL_I2C_FIFOSTAT_RFE   BIT(2)
 #define OWL_I2C_FIFOSTAT_TFF   BIT(5)
 #define OWL_I2C_FIFOSTAT_RFD   GENMASK(15, 8)
 
 /* I2C bus timeout */
-#define OWL_I2C_TIMEOUT                msecs_to_jiffies(4 * 1000)
+#define OWL_I2C_TIMEOUT_MS     (4 * 1000)
+#define OWL_I2C_TIMEOUT                msecs_to_jiffies(OWL_I2C_TIMEOUT_MS)
 
 #define OWL_I2C_MAX_RETRIES    50
 
        writel(OWL_I2C_DIV_FACTOR(val), i2c_dev->base + OWL_I2C_REG_CLKDIV);
 }
 
-static irqreturn_t owl_i2c_interrupt(int irq, void *_dev)
+static void owl_i2c_xfer_data(struct owl_i2c_dev *i2c_dev)
 {
-       struct owl_i2c_dev *i2c_dev = _dev;
        struct i2c_msg *msg = i2c_dev->msg;
        unsigned int stat, fifostat;
 
-       spin_lock(&i2c_dev->lock);
-
        i2c_dev->err = 0;
 
        /* Handle NACK from slave */
                /* Clear NACK error bit by writing "1" */
                owl_i2c_update_reg(i2c_dev->base + OWL_I2C_REG_FIFOSTAT,
                                   OWL_I2C_FIFOSTAT_RNB, true);
-               goto stop;
+               return;
        }
 
        /* Handle bus error */
                /* Clear BUS error bit by writing "1" */
                owl_i2c_update_reg(i2c_dev->base + OWL_I2C_REG_STAT,
                                   OWL_I2C_STAT_BEB, true);
-               goto stop;
+               return;
        }
 
        /* Handle FIFO read */
                               i2c_dev->base + OWL_I2C_REG_TXDAT);
                }
        }
+}
+
+static irqreturn_t owl_i2c_interrupt(int irq, void *_dev)
+{
+       struct owl_i2c_dev *i2c_dev = _dev;
+
+       spin_lock(&i2c_dev->lock);
+
+       owl_i2c_xfer_data(i2c_dev);
 
-stop:
        /* Clear pending interrupts */
        owl_i2c_update_reg(i2c_dev->base + OWL_I2C_REG_STAT,
                           OWL_I2C_STAT_IRQP, true);
        return 0;
 }
 
-static int owl_i2c_master_xfer(struct i2c_adapter *adap, struct i2c_msg *msgs,
-                              int num)
+static int owl_i2c_xfer_common(struct i2c_adapter *adap, struct i2c_msg *msgs,
+                              int num, bool atomic)
 {
        struct owl_i2c_dev *i2c_dev = i2c_get_adapdata(adap);
        struct i2c_msg *msg;
                goto err_exit;
        }
 
-       reinit_completion(&i2c_dev->msg_complete);
+       if (!atomic)
+               reinit_completion(&i2c_dev->msg_complete);
 
-       /* Enable I2C controller interrupt */
+       /* Enable/disable I2C controller interrupt */
        owl_i2c_update_reg(i2c_dev->base + OWL_I2C_REG_CTL,
-                          OWL_I2C_CTL_IRQE, true);
+                          OWL_I2C_CTL_IRQE, !atomic);
 
        /*
         * Select: FIFO enable, Master mode, Stop enable, Data count enable,
 
        spin_unlock_irqrestore(&i2c_dev->lock, flags);
 
-       time_left = wait_for_completion_timeout(&i2c_dev->msg_complete,
-                                               adap->timeout);
+       if (atomic) {
+               /* Wait for Command Execute Completed or NACK Error bits */
+               ret = readl_poll_timeout_atomic(i2c_dev->base + OWL_I2C_REG_FIFOSTAT,
+                                               val, val & (OWL_I2C_FIFOSTAT_CECB |
+                                                           OWL_I2C_FIFOSTAT_RNB),
+                                               10, OWL_I2C_TIMEOUT_MS * 1000);
+       } else {
+               time_left = wait_for_completion_timeout(&i2c_dev->msg_complete,
+                                                       adap->timeout);
+               if (!time_left)
+                       ret = -ETIMEDOUT;
+       }
 
        spin_lock_irqsave(&i2c_dev->lock, flags);
-       if (time_left == 0) {
+
+       if (ret) {
                dev_err(&adap->dev, "Transaction timed out\n");
                /* Send stop condition and release the bus */
                owl_i2c_update_reg(i2c_dev->base + OWL_I2C_REG_CTL,
                                   OWL_I2C_CTL_GBCC_STOP | OWL_I2C_CTL_RB,
                                   true);
-               ret = -ETIMEDOUT;
                goto err_exit;
        }
 
+       if (atomic)
+               owl_i2c_xfer_data(i2c_dev);
+
        ret = i2c_dev->err < 0 ? i2c_dev->err : num;
 
 err_exit:
        return ret;
 }
 
+static int owl_i2c_xfer(struct i2c_adapter *adap, struct i2c_msg *msgs,
+                       int num)
+{
+       return owl_i2c_xfer_common(adap, msgs, num, false);
+}
+
+static int owl_i2c_xfer_atomic(struct i2c_adapter *adap,
+                              struct i2c_msg *msgs, int num)
+{
+       return owl_i2c_xfer_common(adap, msgs, num, true);
+}
+
 static const struct i2c_algorithm owl_i2c_algorithm = {
-       .master_xfer    = owl_i2c_master_xfer,
-       .functionality  = owl_i2c_func,
+       .master_xfer         = owl_i2c_xfer,
+       .master_xfer_atomic  = owl_i2c_xfer_atomic,
+       .functionality       = owl_i2c_func,
 };
 
 static const struct i2c_adapter_quirks owl_i2c_quirks = {