1 files changed, 204 insertions, 175 deletions

diff --git a/drivers/net/wireless/mediatek/mt76/usb.c b/drivers/net/wireless/mediatek/mt76/usb.c
index 4c1abd492405..bbaa1365bbda 100644
--- a/drivers/net/wireless/mediatek/mt76/usb.c
+++ b/drivers/net/wireless/mediatek/mt76/usb.c
@@ -31,8 +31,7 @@ static int __mt76u_vendor_request(struct mt76_dev *dev, u8 req,
 				  u8 req_type, u16 val, u16 offset,
 				  void *buf, size_t len)
 {
-	struct usb_interface *intf = to_usb_interface(dev->dev);
-	struct usb_device *udev = interface_to_usbdev(intf);
+	struct usb_device *udev = to_usb_device(dev->dev);
 	unsigned int pipe;
 	int i, ret;
 
@@ -247,8 +246,7 @@ mt76u_rd_rp(struct mt76_dev *dev, u32 base,
 
 static bool mt76u_check_sg(struct mt76_dev *dev)
 {
-	struct usb_interface *intf = to_usb_interface(dev->dev);
-	struct usb_device *udev = interface_to_usbdev(intf);
+	struct usb_device *udev = to_usb_device(dev->dev);
 
 	return (!disable_usb_sg && udev->bus->sg_tablesize > 0 &&
 		(udev->bus->no_sg_constraint ||
@@ -285,28 +283,24 @@ mt76u_set_endpoints(struct usb_interface *intf,
 }
 
 static int
-mt76u_fill_rx_sg(struct mt76_dev *dev, struct mt76u_buf *buf,
-		 int nsgs, int len, int sglen)
+mt76u_fill_rx_sg(struct mt76_dev *dev, struct mt76_queue *q, struct urb *urb,
+		 int nsgs, gfp_t gfp)
 {
-	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
-	struct urb *urb = buf->urb;
 	int i;
 
-	spin_lock_bh(&q->rx_page_lock);
 	for (i = 0; i < nsgs; i++) {
 		struct page *page;
 		void *data;
 		int offset;
 
-		data = page_frag_alloc(&q->rx_page, len, GFP_ATOMIC);
+		data = page_frag_alloc(&q->rx_page, q->buf_size, gfp);
 		if (!data)
 			break;
 
 		page = virt_to_head_page(data);
 		offset = data - page_address(page);
-		sg_set_page(&urb->sg[i], page, sglen, offset);
+		sg_set_page(&urb->sg[i], page, q->buf_size, offset);
 	}
-	spin_unlock_bh(&q->rx_page_lock);
 
 	if (i < nsgs) {
 		int j;
@@ -317,72 +311,78 @@ mt76u_fill_rx_sg(struct mt76_dev *dev, struct mt76u_buf *buf,
 	}
 
 	urb->num_sgs = max_t(int, i, urb->num_sgs);
-	buf->len = urb->num_sgs * sglen,
+	urb->transfer_buffer_length = urb->num_sgs * q->buf_size,
 	sg_init_marker(urb->sg, urb->num_sgs);
 
 	return i ? : -ENOMEM;
 }
 
 static int
-mt76u_refill_rx(struct mt76_dev *dev, struct mt76_queue *q,
-		struct mt76u_buf *buf, int nsgs, gfp_t gfp)
+mt76u_refill_rx(struct mt76_dev *dev, struct urb *urb, int nsgs, gfp_t gfp)
 {
+	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
+
 	if (dev->usb.sg_en) {
-		return mt76u_fill_rx_sg(dev, buf, nsgs, q->buf_size,
-					SKB_WITH_OVERHEAD(q->buf_size));
+		return mt76u_fill_rx_sg(dev, q, urb, nsgs, gfp);
 	} else {
-		buf->buf = page_frag_alloc(&q->rx_page, q->buf_size, gfp);
-		return buf->buf ? 0 : -ENOMEM;
+		urb->transfer_buffer_length = q->buf_size;
+		urb->transfer_buffer = page_frag_alloc(&q->rx_page,
+						       q->buf_size, gfp);
+		return urb->transfer_buffer ? 0 : -ENOMEM;
 	}
 }
 
 static int
-mt76u_buf_alloc(struct mt76_dev *dev, struct mt76u_buf *buf)
+mt76u_urb_alloc(struct mt76_dev *dev, struct mt76_queue_entry *e)
 {
-	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
+	unsigned int size = sizeof(struct urb);
 
-	buf->len = SKB_WITH_OVERHEAD(q->buf_size);
-	buf->dev = dev;
+	if (dev->usb.sg_en)
+		size += MT_SG_MAX_SIZE * sizeof(struct scatterlist);
 
-	buf->urb = usb_alloc_urb(0, GFP_KERNEL);
-	if (!buf->urb)
+	e->urb = kzalloc(size, GFP_KERNEL);
+	if (!e->urb)
 		return -ENOMEM;
 
-	if (dev->usb.sg_en) {
-		buf->urb->sg = devm_kcalloc(dev->dev, MT_SG_MAX_SIZE,
-					    sizeof(*buf->urb->sg),
-					    GFP_KERNEL);
-		if (!buf->urb->sg)
-			return -ENOMEM;
+	usb_init_urb(e->urb);
 
-		sg_init_table(buf->urb->sg, MT_SG_MAX_SIZE);
-	}
+	if (dev->usb.sg_en)
+		e->urb->sg = (struct scatterlist *)(e->urb + 1);
 
-	return mt76u_refill_rx(dev, q, buf, MT_SG_MAX_SIZE, GFP_KERNEL);
+	return 0;
 }
 
-static void mt76u_buf_free(struct mt76u_buf *buf)
+static int
+mt76u_rx_urb_alloc(struct mt76_dev *dev, struct mt76_queue_entry *e)
+{
+	int err;
+
+	err = mt76u_urb_alloc(dev, e);
+	if (err)
+		return err;
+
+	return mt76u_refill_rx(dev, e->urb, MT_SG_MAX_SIZE, GFP_KERNEL);
+}
+
+static void mt76u_urb_free(struct urb *urb)
 {
-	struct urb *urb = buf->urb;
 	int i;
 
 	for (i = 0; i < urb->num_sgs; i++)
 		skb_free_frag(sg_virt(&urb->sg[i]));
 
-	if (buf->buf)
-		skb_free_frag(buf->buf);
+	if (urb->transfer_buffer)
+		skb_free_frag(urb->transfer_buffer);
 
-	usb_free_urb(buf->urb);
+	usb_free_urb(urb);
 }
 
 static void
 mt76u_fill_bulk_urb(struct mt76_dev *dev, int dir, int index,
-		    struct mt76u_buf *buf, usb_complete_t complete_fn,
+		    struct urb *urb, usb_complete_t complete_fn,
 		    void *context)
 {
-	struct usb_interface *intf = to_usb_interface(dev->dev);
-	struct usb_device *udev = interface_to_usbdev(intf);
-	u8 *data = buf->urb->num_sgs ? NULL : buf->buf;
+	struct usb_device *udev = to_usb_device(dev->dev);
 	unsigned int pipe;
 
 	if (dir == USB_DIR_IN)
@@ -390,37 +390,28 @@ mt76u_fill_bulk_urb(struct mt76_dev *dev, int dir, int index,
 	else
 		pipe = usb_sndbulkpipe(udev, dev->usb.out_ep[index]);
 
-	usb_fill_bulk_urb(buf->urb, udev, pipe, data, buf->len,
-			  complete_fn, context);
+	urb->dev = udev;
+	urb->pipe = pipe;
+	urb->complete = complete_fn;
+	urb->context = context;
 }
 
-static int
-mt76u_submit_buf(struct mt76_dev *dev, int dir, int index,
-		 struct mt76u_buf *buf, gfp_t gfp,
-		 usb_complete_t complete_fn, void *context)
+static inline struct urb *
+mt76u_get_next_rx_entry(struct mt76_dev *dev)
 {
-	mt76u_fill_bulk_urb(dev, dir, index, buf, complete_fn,
-			    context);
-	trace_submit_urb(dev, buf->urb);
-
-	return usb_submit_urb(buf->urb, gfp);
-}
-
-static inline struct mt76u_buf
-*mt76u_get_next_rx_entry(struct mt76_queue *q)
-{
-	struct mt76u_buf *buf = NULL;
+	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
+	struct urb *urb = NULL;
 	unsigned long flags;
 
 	spin_lock_irqsave(&q->lock, flags);
 	if (q->queued > 0) {
-		buf = &q->entry[q->head].ubuf;
+		urb = q->entry[q->head].urb;
 		q->head = (q->head + 1) % q->ndesc;
 		q->queued--;
 	}
 	spin_unlock_irqrestore(&q->lock, flags);
 
-	return buf;
+	return urb;
 }
 
 static int mt76u_get_rx_entry_len(u8 *data, u32 data_len)
@@ -439,12 +430,12 @@ static int mt76u_get_rx_entry_len(u8 *data, u32 data_len)
 }
 
 static int
-mt76u_process_rx_entry(struct mt76_dev *dev, struct mt76u_buf *buf)
+mt76u_process_rx_entry(struct mt76_dev *dev, struct urb *urb)
 {
 	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
-	struct urb *urb = buf->urb;
-	u8 *data = urb->num_sgs ? sg_virt(&urb->sg[0]) : buf->buf;
-	int data_len, len, nsgs = 1;
+	u8 *data = urb->num_sgs ? sg_virt(&urb->sg[0]) : urb->transfer_buffer;
+	int data_len = urb->num_sgs ? urb->sg[0].length : urb->actual_length;
+	int len, nsgs = 1;
 	struct sk_buff *skb;
 
 	if (!test_bit(MT76_STATE_INITIALIZED, &dev->state))
@@ -454,10 +445,11 @@ mt76u_process_rx_entry(struct mt76_dev *dev, struct mt76u_buf *buf)
 	if (len < 0)
 		return 0;
 
-	data_len = urb->num_sgs ? urb->sg[0].length : buf->len;
 	data_len = min_t(int, len, data_len - MT_DMA_HDR_LEN);
-	if (MT_DMA_HDR_LEN + data_len > SKB_WITH_OVERHEAD(q->buf_size))
+	if (MT_DMA_HDR_LEN + data_len > SKB_WITH_OVERHEAD(q->buf_size)) {
+		dev_err_ratelimited(dev->dev, "rx data too big %d\n", data_len);
 		return 0;
+	}
 
 	skb = build_skb(data, q->buf_size);
 	if (!skb)
@@ -503,7 +495,7 @@ static void mt76u_complete_rx(struct urb *urb)
 	}
 
 	spin_lock_irqsave(&q->lock, flags);
-	if (WARN_ONCE(q->entry[q->tail].ubuf.urb != urb, "rx urb mismatch"))
+	if (WARN_ONCE(q->entry[q->tail].urb != urb, "rx urb mismatch"))
 		goto out;
 
 	q->tail = (q->tail + 1) % q->ndesc;
@@ -513,37 +505,43 @@ out:
 	spin_unlock_irqrestore(&q->lock, flags);
 }
 
+static int
+mt76u_submit_rx_buf(struct mt76_dev *dev, struct urb *urb)
+{
+	mt76u_fill_bulk_urb(dev, USB_DIR_IN, MT_EP_IN_PKT_RX, urb,
+			    mt76u_complete_rx, dev);
+	trace_submit_urb(dev, urb);
+
+	return usb_submit_urb(urb, GFP_ATOMIC);
+}
+
 static void mt76u_rx_tasklet(unsigned long data)
 {
 	struct mt76_dev *dev = (struct mt76_dev *)data;
-	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
-	struct mt76u_buf *buf;
+	struct urb *urb;
 	int err, count;
 
 	rcu_read_lock();
 
 	while (true) {
-		buf = mt76u_get_next_rx_entry(q);
-		if (!buf)
+		urb = mt76u_get_next_rx_entry(dev);
+		if (!urb)
 			break;
 
-		count = mt76u_process_rx_entry(dev, buf);
+		count = mt76u_process_rx_entry(dev, urb);
 		if (count > 0) {
-			err = mt76u_refill_rx(dev, q, buf, count,
-					      GFP_ATOMIC);
+			err = mt76u_refill_rx(dev, urb, count, GFP_ATOMIC);
 			if (err < 0)
 				break;
 		}
-		mt76u_submit_buf(dev, USB_DIR_IN, MT_EP_IN_PKT_RX,
-				 buf, GFP_ATOMIC,
-				 mt76u_complete_rx, dev);
+		mt76u_submit_rx_buf(dev, urb);
 	}
 	mt76_rx_poll_complete(dev, MT_RXQ_MAIN, NULL);
 
 	rcu_read_unlock();
 }
 
-int mt76u_submit_rx_buffers(struct mt76_dev *dev)
+static int mt76u_submit_rx_buffers(struct mt76_dev *dev)
 {
 	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
 	unsigned long flags;
@@ -551,9 +549,7 @@ int mt76u_submit_rx_buffers(struct mt76_dev *dev)
 
 	spin_lock_irqsave(&q->lock, flags);
 	for (i = 0; i < q->ndesc; i++) {
-		err = mt76u_submit_buf(dev, USB_DIR_IN, MT_EP_IN_PKT_RX,
-				       &q->entry[i].ubuf, GFP_ATOMIC,
-				       mt76u_complete_rx, dev);
+		err = mt76u_submit_rx_buf(dev, q->entry[i].urb);
 		if (err < 0)
 			break;
 	}
@@ -563,7 +559,6 @@ int mt76u_submit_rx_buffers(struct mt76_dev *dev)
 
 	return err;
 }
-EXPORT_SYMBOL_GPL(mt76u_submit_rx_buffers);
 
 static int mt76u_alloc_rx(struct mt76_dev *dev)
 {
@@ -575,7 +570,6 @@ static int mt76u_alloc_rx(struct mt76_dev *dev)
 	if (!usb->mcu.data)
 		return -ENOMEM;
 
-	spin_lock_init(&q->rx_page_lock);
 	spin_lock_init(&q->lock);
 	q->entry = devm_kcalloc(dev->dev,
 				MT_NUM_RX_ENTRIES, sizeof(*q->entry),
@@ -586,7 +580,7 @@ static int mt76u_alloc_rx(struct mt76_dev *dev)
 	q->buf_size = dev->usb.sg_en ? MT_RX_BUF_SIZE : PAGE_SIZE;
 	q->ndesc = MT_NUM_RX_ENTRIES;
 	for (i = 0; i < q->ndesc; i++) {
-		err = mt76u_buf_alloc(dev, &q->entry[i].ubuf);
+		err = mt76u_rx_urb_alloc(dev, &q->entry[i]);
 		if (err < 0)
 			return err;
 	}
@@ -601,60 +595,76 @@ static void mt76u_free_rx(struct mt76_dev *dev)
 	int i;
 
 	for (i = 0; i < q->ndesc; i++)
-		mt76u_buf_free(&q->entry[i].ubuf);
+		mt76u_urb_free(q->entry[i].urb);
 
-	spin_lock_bh(&q->rx_page_lock);
 	if (!q->rx_page.va)
-		goto out;
+		return;
 
 	page = virt_to_page(q->rx_page.va);
 	__page_frag_cache_drain(page, q->rx_page.pagecnt_bias);
 	memset(&q->rx_page, 0, sizeof(q->rx_page));
-out:
-	spin_unlock_bh(&q->rx_page_lock);
 }
 
-static void mt76u_stop_rx(struct mt76_dev *dev)
+void mt76u_stop_rx(struct mt76_dev *dev)
 {
 	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
 	int i;
 
 	for (i = 0; i < q->ndesc; i++)
-		usb_kill_urb(q->entry[i].ubuf.urb);
+		usb_poison_urb(q->entry[i].urb);
+
+	tasklet_kill(&dev->usb.rx_tasklet);
 }
+EXPORT_SYMBOL_GPL(mt76u_stop_rx);
+
+int mt76u_resume_rx(struct mt76_dev *dev)
+{
+	struct mt76_queue *q = &dev->q_rx[MT_RXQ_MAIN];
+	int i;
+
+	for (i = 0; i < q->ndesc; i++)
+		usb_unpoison_urb(q->entry[i].urb);
+
+	return mt76u_submit_rx_buffers(dev);
+}
+EXPORT_SYMBOL_GPL(mt76u_resume_rx);
 
 static void mt76u_tx_tasklet(unsigned long data)
 {
 	struct mt76_dev *dev = (struct mt76_dev *)data;
 	struct mt76_queue_entry entry;
-	struct mt76u_buf *buf;
+	struct mt76_sw_queue *sq;
 	struct mt76_queue *q;
 	bool wake;
 	int i;
 
 	for (i = 0; i < IEEE80211_NUM_ACS; i++) {
-		q = &dev->q_tx[i];
+		u32 n_dequeued = 0, n_sw_dequeued = 0;
 
-		spin_lock_bh(&q->lock);
-		while (true) {
-			buf = &q->entry[q->head].ubuf;
-			if (!buf->done || !q->queued)
+		sq = &dev->q_tx[i];
+		q = sq->q;
+
+		while (q->queued > n_dequeued) {
+			if (!q->entry[q->head].done)
 				break;
 
 			if (q->entry[q->head].schedule) {
 				q->entry[q->head].schedule = false;
-				q->swq_queued--;
+				n_sw_dequeued++;
 			}
 
 			entry = q->entry[q->head];
+			q->entry[q->head].done = false;
 			q->head = (q->head + 1) % q->ndesc;
-			q->queued--;
+			n_dequeued++;
 
-			spin_unlock_bh(&q->lock);
-			dev->drv->tx_complete_skb(dev, q, &entry, false);
-			spin_lock_bh(&q->lock);
+			dev->drv->tx_complete_skb(dev, i, &entry);
 		}
-		mt76_txq_schedule(dev, q);
+
+		spin_lock_bh(&q->lock);
+
+		sq->swq_queued -= n_sw_dequeued;
+		q->queued -= n_dequeued;
 
 		wake = q->stopped && q->queued < q->ndesc - 8;
 		if (wake)
@@ -665,6 +675,8 @@ static void mt76u_tx_tasklet(unsigned long data)
 
 		spin_unlock_bh(&q->lock);
 
+		mt76_txq_schedule(dev, i);
+
 		if (!test_and_set_bit(MT76_READING_STATS, &dev->state))
 			ieee80211_queue_delayed_work(dev->hw,
 						     &dev->usb.stat_work,
@@ -703,34 +715,43 @@ static void mt76u_tx_status_data(struct work_struct *work)
 
 static void mt76u_complete_tx(struct urb *urb)
 {
-	struct mt76u_buf *buf = urb->context;
-	struct mt76_dev *dev = buf->dev;
+	struct mt76_dev *dev = dev_get_drvdata(&urb->dev->dev);
+	struct mt76_queue_entry *e = urb->context;
 
 	if (mt76u_urb_error(urb))
 		dev_err(dev->dev, "tx urb failed: %d\n", urb->status);
-	buf->done = true;
+	e->done = true;
 
-	tasklet_schedule(&dev->usb.tx_tasklet);
+	tasklet_schedule(&dev->tx_tasklet);
 }
 
 static int
-mt76u_tx_build_sg(struct mt76_dev *dev, struct sk_buff *skb,
-		  struct urb *urb)
+mt76u_tx_setup_buffers(struct mt76_dev *dev, struct sk_buff *skb,
+		       struct urb *urb)
 {
-	if (!dev->usb.sg_en)
-		return 0;
+	urb->transfer_buffer_length = skb->len;
 
-	sg_init_table(urb->sg, MT_SG_MAX_SIZE);
-	urb->num_sgs = skb_to_sgvec(skb, urb->sg, 0, skb->len);
-	return urb->num_sgs;
+	if (!dev->usb.sg_en) {
+		urb->transfer_buffer = skb->data;
+		return 0;
+	} else {
+		sg_init_table(urb->sg, MT_SG_MAX_SIZE);
+		urb->num_sgs = skb_to_sgvec(skb, urb->sg, 0, skb->len);
+		if (urb->num_sgs == 0)
+			return -ENOMEM;
+		return urb->num_sgs;
+	}
 }
 
 static int
-mt76u_tx_queue_skb(struct mt76_dev *dev, struct mt76_queue *q,
+mt76u_tx_queue_skb(struct mt76_dev *dev, enum mt76_txq_id qid,
 		   struct sk_buff *skb, struct mt76_wcid *wcid,
 		   struct ieee80211_sta *sta)
 {
-	struct mt76u_buf *buf;
+	struct mt76_queue *q = dev->q_tx[qid].q;
+	struct mt76_tx_info tx_info = {
+		.skb = skb,
+	};
 	u16 idx = q->tail;
 	int err;
 
@@ -738,24 +759,20 @@ mt76u_tx_queue_skb(struct mt76_dev *dev, struct mt76_queue *q,
 		return -ENOSPC;
 
 	skb->prev = skb->next = NULL;
-	err = dev->drv->tx_prepare_skb(dev, NULL, skb, q, wcid, sta, NULL);
+	err = dev->drv->tx_prepare_skb(dev, NULL, qid, wcid, sta, &tx_info);
 	if (err < 0)
 		return err;
 
-	buf = &q->entry[idx].ubuf;
-	buf->buf = skb->data;
-	buf->len = skb->len;
-	buf->done = false;
-
-	err = mt76u_tx_build_sg(dev, skb, buf->urb);
+	err = mt76u_tx_setup_buffers(dev, tx_info.skb, q->entry[idx].urb);
 	if (err < 0)
 		return err;
 
 	mt76u_fill_bulk_urb(dev, USB_DIR_OUT, q2ep(q->hw_idx),
-			    buf, mt76u_complete_tx, buf);
+			    q->entry[idx].urb, mt76u_complete_tx,
+			    &q->entry[idx]);
 
 	q->tail = (q->tail + 1) % q->ndesc;
-	q->entry[idx].skb = skb;
+	q->entry[idx].skb = tx_info.skb;
 	q->queued++;
 
 	return idx;
@@ -763,14 +780,14 @@ mt76u_tx_queue_skb(struct mt76_dev *dev, struct mt76_queue *q,
 
 static void mt76u_tx_kick(struct mt76_dev *dev, struct mt76_queue *q)
 {
-	struct mt76u_buf *buf;
+	struct urb *urb;
 	int err;
 
 	while (q->first != q->tail) {
-		buf = &q->entry[q->first].ubuf;
+		urb = q->entry[q->first].urb;
 
-		trace_submit_urb(dev, buf->urb);
-		err = usb_submit_urb(buf->urb, GFP_ATOMIC);
+		trace_submit_urb(dev, urb);
+		err = usb_submit_urb(urb, GFP_ATOMIC);
 		if (err < 0) {
 			if (err == -ENODEV)
 				set_bit(MT76_REMOVED, &dev->state);
@@ -785,15 +802,24 @@ static void mt76u_tx_kick(struct mt76_dev *dev, struct mt76_queue *q)
 
 static int mt76u_alloc_tx(struct mt76_dev *dev)
 {
-	struct mt76u_buf *buf;
 	struct mt76_queue *q;
-	int i, j;
+	int i, j, err;
+
+	for (i = 0; i <= MT_TXQ_PSD; i++) {
+		INIT_LIST_HEAD(&dev->q_tx[i].swq);
+
+		if (i >= IEEE80211_NUM_ACS) {
+			dev->q_tx[i].q = dev->q_tx[0].q;
+			continue;
+		}
+
+		q = devm_kzalloc(dev->dev, sizeof(*q), GFP_KERNEL);
+		if (!q)
+			return -ENOMEM;
 
-	for (i = 0; i < IEEE80211_NUM_ACS; i++) {
-		q = &dev->q_tx[i];
 		spin_lock_init(&q->lock);
-		INIT_LIST_HEAD(&q->swq);
 		q->hw_idx = mt76_ac_to_hwq(i);
+		dev->q_tx[i].q = q;
 
 		q->entry = devm_kcalloc(dev->dev,
 					MT_NUM_TX_ENTRIES, sizeof(*q->entry),
@@ -803,22 +829,9 @@ static int mt76u_alloc_tx(struct mt76_dev *dev)
 
 		q->ndesc = MT_NUM_TX_ENTRIES;
 		for (j = 0; j < q->ndesc; j++) {
-			buf = &q->entry[j].ubuf;
-			buf->dev = dev;
-
-			buf->urb = usb_alloc_urb(0, GFP_KERNEL);
-			if (!buf->urb)
-				return -ENOMEM;
-
-			if (dev->usb.sg_en) {
-				size_t size = MT_SG_MAX_SIZE *
-					      sizeof(struct scatterlist);
-
-				buf->urb->sg = devm_kzalloc(dev->dev, size,
-							    GFP_KERNEL);
-				if (!buf->urb->sg)
-					return -ENOMEM;
-			}
+			err = mt76u_urb_alloc(dev, &q->entry[j]);
+			if (err < 0)
+				return err;
 		}
 	}
 	return 0;
@@ -830,44 +843,60 @@ static void mt76u_free_tx(struct mt76_dev *dev)
 	int i, j;
 
 	for (i = 0; i < IEEE80211_NUM_ACS; i++) {
-		q = &dev->q_tx[i];
+		q = dev->q_tx[i].q;
 		for (j = 0; j < q->ndesc; j++)
-			usb_free_urb(q->entry[j].ubuf.urb);
+			usb_free_urb(q->entry[j].urb);
 	}
 }
 
-static void mt76u_stop_tx(struct mt76_dev *dev)
+void mt76u_stop_tx(struct mt76_dev *dev)
 {
+	struct mt76_queue_entry entry;
 	struct mt76_queue *q;
-	int i, j;
+	int i, j, ret;
 
-	for (i = 0; i < IEEE80211_NUM_ACS; i++) {
-		q = &dev->q_tx[i];
-		for (j = 0; j < q->ndesc; j++)
-			usb_kill_urb(q->entry[j].ubuf.urb);
-	}
-}
+	ret = wait_event_timeout(dev->tx_wait, !mt76_has_tx_pending(dev), HZ/5);
+	if (!ret) {
+		dev_err(dev->dev, "timed out waiting for pending tx\n");
 
-void mt76u_stop_queues(struct mt76_dev *dev)
-{
-	tasklet_disable(&dev->usb.rx_tasklet);
-	tasklet_disable(&dev->usb.tx_tasklet);
+		for (i = 0; i < IEEE80211_NUM_ACS; i++) {
+			q = dev->q_tx[i].q;
+			for (j = 0; j < q->ndesc; j++)
+				usb_kill_urb(q->entry[j].urb);
+		}
 
-	mt76u_stop_rx(dev);
-	mt76u_stop_tx(dev);
-}
-EXPORT_SYMBOL_GPL(mt76u_stop_queues);
+		tasklet_kill(&dev->tx_tasklet);
+
+		/* On device removal we maight queue skb's, but mt76u_tx_kick()
+		 * will fail to submit urb, cleanup those skb's manually.
+		 */
+		for (i = 0; i < IEEE80211_NUM_ACS; i++) {
+			q = dev->q_tx[i].q;
+
+			/* Assure we are in sync with killed tasklet. */
+			spin_lock_bh(&q->lock);
+			while (q->queued) {
+				entry = q->entry[q->head];
+				q->head = (q->head + 1) % q->ndesc;
+				q->queued--;
+
+				dev->drv->tx_complete_skb(dev, i, &entry);
+			}
+			spin_unlock_bh(&q->lock);
+		}
+	}
 
-void mt76u_stop_stat_wk(struct mt76_dev *dev)
-{
 	cancel_delayed_work_sync(&dev->usb.stat_work);
 	clear_bit(MT76_READING_STATS, &dev->state);
+
+	mt76_tx_status_check(dev, NULL, true);
 }
-EXPORT_SYMBOL_GPL(mt76u_stop_stat_wk);
+EXPORT_SYMBOL_GPL(mt76u_stop_tx);
 
 void mt76u_queues_deinit(struct mt76_dev *dev)
 {
-	mt76u_stop_queues(dev);
+	mt76u_stop_rx(dev);
+	mt76u_stop_tx(dev);
 
 	mt76u_free_rx(dev);
 	mt76u_free_tx(dev);
@@ -906,7 +935,7 @@ int mt76u_init(struct mt76_dev *dev,
 	struct mt76_usb *usb = &dev->usb;
 
 	tasklet_init(&usb->rx_tasklet, mt76u_rx_tasklet, (unsigned long)dev);
-	tasklet_init(&usb->tx_tasklet, mt76u_tx_tasklet, (unsigned long)dev);
+	tasklet_init(&dev->tx_tasklet, mt76u_tx_tasklet, (unsigned long)dev);
 	INIT_DELAYED_WORK(&usb->stat_work, mt76u_tx_status_data);
 	skb_queue_head_init(&dev->rx_skb[MT_RXQ_MAIN]);