dma_addr_t rx_dma;
        enum geni_se_xfer_mode mode;
        unsigned long time_left = XFER_TIMEOUT;
+       void *dma_buf;
 
        gi2c->cur = msg;
-       mode = msg->len > 32 ? GENI_SE_DMA : GENI_SE_FIFO;
+       mode = GENI_SE_FIFO;
+       dma_buf = i2c_get_dma_safe_msg_buf(msg, 32);
+       if (dma_buf)
+               mode = GENI_SE_DMA;
+
        geni_se_select_mode(&gi2c->se, mode);
        writel_relaxed(msg->len, gi2c->se.base + SE_I2C_RX_TRANS_LEN);
        geni_se_setup_m_cmd(&gi2c->se, I2C_READ, m_param);
        if (mode == GENI_SE_DMA) {
                int ret;
 
-               ret = geni_se_rx_dma_prep(&gi2c->se, msg->buf, msg->len,
+               ret = geni_se_rx_dma_prep(&gi2c->se, dma_buf, msg->len,
                                                                &rx_dma);
                if (ret) {
                        mode = GENI_SE_FIFO;
                        geni_se_select_mode(&gi2c->se, mode);
+                       i2c_put_dma_safe_msg_buf(dma_buf, msg, false);
                }
        }
 
                if (gi2c->err)
                        geni_i2c_rx_fsm_rst(gi2c);
                geni_se_rx_dma_unprep(&gi2c->se, rx_dma, msg->len);
+               i2c_put_dma_safe_msg_buf(dma_buf, msg, !gi2c->err);
        }
        return gi2c->err;
 }
        dma_addr_t tx_dma;
        enum geni_se_xfer_mode mode;
        unsigned long time_left;
+       void *dma_buf;
 
        gi2c->cur = msg;
-       mode = msg->len > 32 ? GENI_SE_DMA : GENI_SE_FIFO;
+       mode = GENI_SE_FIFO;
+       dma_buf = i2c_get_dma_safe_msg_buf(msg, 32);
+       if (dma_buf)
+               mode = GENI_SE_DMA;
+
        geni_se_select_mode(&gi2c->se, mode);
        writel_relaxed(msg->len, gi2c->se.base + SE_I2C_TX_TRANS_LEN);
        geni_se_setup_m_cmd(&gi2c->se, I2C_WRITE, m_param);
        if (mode == GENI_SE_DMA) {
                int ret;
 
-               ret = geni_se_tx_dma_prep(&gi2c->se, msg->buf, msg->len,
+               ret = geni_se_tx_dma_prep(&gi2c->se, dma_buf, msg->len,
                                                                &tx_dma);
                if (ret) {
                        mode = GENI_SE_FIFO;
                        geni_se_select_mode(&gi2c->se, mode);
+                       i2c_put_dma_safe_msg_buf(dma_buf, msg, false);
                }
        }
 
                if (gi2c->err)
                        geni_i2c_tx_fsm_rst(gi2c);
                geni_se_tx_dma_unprep(&gi2c->se, tx_dma, msg->len);
+               i2c_put_dma_safe_msg_buf(dma_buf, msg, !gi2c->err);
        }
        return gi2c->err;
 }