#include "mt76x02.h"
 
 static int
-mt76x02_init_tx_queue(struct mt76_dev *dev, struct mt76_queue *q,
+mt76x02_init_tx_queue(struct mt76x02_dev *dev, struct mt76_queue *q,
                      int idx, int n_desc)
 {
        int ret;
 
-       q->regs = dev->mmio.regs + MT_TX_RING_BASE + idx * MT_RING_SIZE;
+       q->regs = dev->mt76.mmio.regs + MT_TX_RING_BASE + idx * MT_RING_SIZE;
        q->ndesc = n_desc;
        q->hw_idx = idx;
 
-       ret = __mt76_queue_alloc(dev, q);
+       ret = mt76_queue_alloc(dev, q);
        if (ret)
                return ret;
 
 }
 
 static int
-mt76x02_init_rx_queue(struct mt76_dev *dev, struct mt76_queue *q,
+mt76x02_init_rx_queue(struct mt76x02_dev *dev, struct mt76_queue *q,
                      int idx, int n_desc, int bufsize)
 {
        int ret;
 
-       q->regs = dev->mmio.regs + MT_RX_RING_BASE + idx * MT_RING_SIZE;
+       q->regs = dev->mt76.mmio.regs + MT_RX_RING_BASE + idx * MT_RING_SIZE;
        q->ndesc = n_desc;
        q->buf_size = bufsize;
 
-       ret = __mt76_queue_alloc(dev, q);
+       ret = mt76_queue_alloc(dev, q);
        if (ret)
                return ret;
 
        return 0;
 }
 
-int mt76x02_dma_init(struct mt76_dev *dev)
+int mt76x02_dma_init(struct mt76x02_dev *dev)
 {
        struct mt76_txwi_cache __maybe_unused *t;
        struct mt76_queue *q;
        BUILD_BUG_ON(sizeof(t->txwi) < sizeof(struct mt76x02_txwi));
        BUILD_BUG_ON(sizeof(struct mt76x02_rxwi) > MT_RX_HEADROOM);
 
-       mt76_dma_attach(dev);
-       __mt76_wr(dev, MT_WPDMA_RST_IDX, ~0);
+       mt76_dma_attach(&dev->mt76);
+       mt76_wr(dev, MT_WPDMA_RST_IDX, ~0);
 
        for (i = 0; i < IEEE80211_NUM_ACS; i++) {
-               ret = mt76x02_init_tx_queue(dev, &dev->q_tx[i],
+               ret = mt76x02_init_tx_queue(dev, &dev->mt76.q_tx[i],
                                            mt76_ac_to_hwq(i),
                                            MT_TX_RING_SIZE);
                if (ret)
                        return ret;
        }
 
-       ret = mt76x02_init_tx_queue(dev, &dev->q_tx[MT_TXQ_PSD],
+       ret = mt76x02_init_tx_queue(dev, &dev->mt76.q_tx[MT_TXQ_PSD],
                                    MT_TX_HW_QUEUE_MGMT, MT_TX_RING_SIZE);
        if (ret)
                return ret;
 
-       ret = mt76x02_init_tx_queue(dev, &dev->q_tx[MT_TXQ_MCU],
+       ret = mt76x02_init_tx_queue(dev, &dev->mt76.q_tx[MT_TXQ_MCU],
                                    MT_TX_HW_QUEUE_MCU, MT_MCU_RING_SIZE);
        if (ret)
                return ret;
 
-       ret = mt76x02_init_rx_queue(dev, &dev->q_rx[MT_RXQ_MCU], 1,
+       ret = mt76x02_init_rx_queue(dev, &dev->mt76.q_rx[MT_RXQ_MCU], 1,
                                    MT_MCU_RING_SIZE, MT_RX_BUF_SIZE);
        if (ret)
                return ret;
 
-       q = &dev->q_rx[MT_RXQ_MAIN];
+       q = &dev->mt76.q_rx[MT_RXQ_MAIN];
        q->buf_offset = MT_RX_HEADROOM - sizeof(struct mt76x02_rxwi);
        ret = mt76x02_init_rx_queue(dev, q, 0, MT76X02_RX_RING_SIZE,
                                    MT_RX_BUF_SIZE);
        if (ret)
                return ret;
 
-       return __mt76_init_queues(dev);
+       return mt76_init_queues(dev);
 }
 EXPORT_SYMBOL_GPL(mt76x02_dma_init);
 
-void mt76x02_set_irq_mask(struct mt76_dev *dev, u32 clear, u32 set)
+void mt76x02_set_irq_mask(struct mt76x02_dev *dev, u32 clear, u32 set)
 {
        unsigned long flags;
 
-       spin_lock_irqsave(&dev->mmio.irq_lock, flags);
-       dev->mmio.irqmask &= ~clear;
-       dev->mmio.irqmask |= set;
-       __mt76_wr(dev, MT_INT_MASK_CSR, dev->mmio.irqmask);
-       spin_unlock_irqrestore(&dev->mmio.irq_lock, flags);
+       spin_lock_irqsave(&dev->mt76.mmio.irq_lock, flags);
+       dev->mt76.mmio.irqmask &= ~clear;
+       dev->mt76.mmio.irqmask |= set;
+       mt76_wr(dev, MT_INT_MASK_CSR, dev->mt76.mmio.irqmask);
+       spin_unlock_irqrestore(&dev->mt76.mmio.irq_lock, flags);
 }
 EXPORT_SYMBOL_GPL(mt76x02_set_irq_mask);
 
-void mt76x02_dma_enable(struct mt76_dev *dev)
+static void mt76x02_dma_enable(struct mt76x02_dev *dev)
 {
        u32 val;
 
-       __mt76_wr(dev, MT_MAC_SYS_CTRL, MT_MAC_SYS_CTRL_ENABLE_TX);
-       mt76x02_wait_for_wpdma(dev, 1000);
+       mt76_wr(dev, MT_MAC_SYS_CTRL, MT_MAC_SYS_CTRL_ENABLE_TX);
+       mt76x02_wait_for_wpdma(&dev->mt76, 1000);
        usleep_range(50, 100);
 
        val = FIELD_PREP(MT_WPDMA_GLO_CFG_DMA_BURST_SIZE, 3) |
              MT_WPDMA_GLO_CFG_TX_DMA_EN |
              MT_WPDMA_GLO_CFG_RX_DMA_EN;
-       __mt76_set(dev, MT_WPDMA_GLO_CFG, val);
-       __mt76_clear(dev, MT_WPDMA_GLO_CFG,
-                    MT_WPDMA_GLO_CFG_TX_WRITEBACK_DONE);
+       mt76_set(dev, MT_WPDMA_GLO_CFG, val);
+       mt76_clear(dev, MT_WPDMA_GLO_CFG,
+                  MT_WPDMA_GLO_CFG_TX_WRITEBACK_DONE);
 }
 EXPORT_SYMBOL_GPL(mt76x02_dma_enable);
 
-void mt76x02_dma_disable(struct mt76_dev *dev)
+void mt76x02_dma_disable(struct mt76x02_dev *dev)
 {
-       u32 val = __mt76_rr(dev, MT_WPDMA_GLO_CFG);
+       u32 val = mt76_rr(dev, MT_WPDMA_GLO_CFG);
 
        val &= MT_WPDMA_GLO_CFG_DMA_BURST_SIZE |
               MT_WPDMA_GLO_CFG_BIG_ENDIAN |
               MT_WPDMA_GLO_CFG_HDR_SEG_LEN;
        val |= MT_WPDMA_GLO_CFG_TX_WRITEBACK_DONE;
-       __mt76_wr(dev, MT_WPDMA_GLO_CFG, val);
+       mt76_wr(dev, MT_WPDMA_GLO_CFG, val);
 }
 EXPORT_SYMBOL_GPL(mt76x02_dma_disable);
 
-void mt76x02_mac_start(struct mt76_dev *dev)
+void mt76x02_mac_start(struct mt76x02_dev *dev)
 {
        mt76x02_dma_enable(dev);
-       __mt76_wr(dev, MT_RX_FILTR_CFG, dev->rxfilter);
-       __mt76_wr(dev, MT_MAC_SYS_CTRL,
-                 MT_MAC_SYS_CTRL_ENABLE_TX |
-                 MT_MAC_SYS_CTRL_ENABLE_RX);
+       mt76_wr(dev, MT_RX_FILTR_CFG, dev->mt76.rxfilter);
+       mt76_wr(dev, MT_MAC_SYS_CTRL,
+               MT_MAC_SYS_CTRL_ENABLE_TX |
+               MT_MAC_SYS_CTRL_ENABLE_RX);
        mt76x02_irq_enable(dev,
                           MT_INT_RX_DONE_ALL | MT_INT_TX_DONE_ALL |
                           MT_INT_TX_STAT);
 
 
 void mt76x2_rx_poll_complete(struct mt76_dev *mdev, enum mt76_rxq_id q)
 {
-       mt76x02_irq_enable(mdev, MT_INT_RX_DONE(q));
+       struct mt76x02_dev *dev;
+
+       dev = container_of(mdev, struct mt76x02_dev, mt76);
+       mt76x02_irq_enable(dev, MT_INT_RX_DONE(q));
 }
 
 irqreturn_t mt76x2_irq_handler(int irq, void *dev_instance)
        intr &= dev->mt76.mmio.irqmask;
 
        if (intr & MT_INT_TX_DONE_ALL) {
-               mt76x02_irq_disable(&dev->mt76, MT_INT_TX_DONE_ALL);
+               mt76x02_irq_disable(dev, MT_INT_TX_DONE_ALL);
                tasklet_schedule(&dev->tx_tasklet);
        }
 
        if (intr & MT_INT_RX_DONE(0)) {
-               mt76x02_irq_disable(&dev->mt76, MT_INT_RX_DONE(0));
+               mt76x02_irq_disable(dev, MT_INT_RX_DONE(0));
                napi_schedule(&dev->mt76.napi[0]);
        }
 
        if (intr & MT_INT_RX_DONE(1)) {
-               mt76x02_irq_disable(&dev->mt76, MT_INT_RX_DONE(1));
+               mt76x02_irq_disable(dev, MT_INT_RX_DONE(1));
                napi_schedule(&dev->mt76.napi[1]);
        }
 
        }
 
        if (intr & MT_INT_GPTIMER) {
-               mt76x02_irq_disable(&dev->mt76, MT_INT_GPTIMER);
+               mt76x02_irq_disable(dev, MT_INT_GPTIMER);
                tasklet_schedule(&dev->dfs_pd.dfs_tasklet);
        }