usb: dwc3: gadget: Revise preparation for extra TRBs

index cb19d682bbac..44351078f833 100644
@@ -1098,6 +1098,8 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 	struct scatterlist *s;
 	int		i;
 	unsigned int length = req->request.length;
+	unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
+	unsigned int rem = length % maxp;
 	unsigned int remaining = req->request.num_mapped_sgs
 		- req->num_queued_sgs;
 

@@ -1109,8 +1111,6 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		length -= sg_dma_len(s);
 
 	for_each_sg(sg, s, remaining, i) {
-		unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
-		unsigned int rem = length % maxp;
 		unsigned int trb_length;
 		unsigned int chain = true;
 

index 44351078f833..6017263929f3 100644
@@ -2738,6 +2738,11 @@ static int dwc3_gadget_ep_cleanup_completed_request(struct dwc3_ep *dep,
 		ret = dwc3_gadget_ep_reclaim_trb_linear(dep, req, event,
 				status);
 
+	req->request.actual = req->request.length - req->remaining;
+
+	if (!dwc3_gadget_ep_request_completed(req))
+		goto out;
+
 	if (req->needs_extra_trb) {
 		unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
 

@@ -2753,11 +2758,6 @@ static int dwc3_gadget_ep_cleanup_completed_request(struct dwc3_ep *dep,
 		req->needs_extra_trb = false;
 	}
 
-	req->request.actual = req->request.length - req->remaining;
-
-	if (!dwc3_gadget_ep_request_completed(req))
-		goto out;
-
 	dwc3_gadget_giveback(dep, req, status);
 
 out:

index 6017263929f3..c0309ad34d28 100644
@@ -1060,10 +1060,11 @@ static void __dwc3_prepare_one_trb(struct dwc3_ep *dep, struct dwc3_trb *trb,
  * @trb_length: buffer size of the TRB
  * @chain: should this TRB be chained to the next?
  * @node: only for isochronous endpoints. First TRB needs different type.
+ * @use_bounce_buffer: set to use bounce buffer
  */
 static void dwc3_prepare_one_trb(struct dwc3_ep *dep,
 		struct dwc3_request *req, unsigned int trb_length,
-		unsigned int chain, unsigned int node)
+		unsigned int chain, unsigned int node, bool use_bounce_buffer)
 {
 	struct dwc3_trb		*trb;
 	dma_addr_t		dma;

@@ -1072,7 +1073,9 @@ static void dwc3_prepare_one_trb(struct dwc3_ep *dep,
 	unsigned int		no_interrupt = req->request.no_interrupt;
 	unsigned int		is_last = req->request.is_last;
 
-	if (req->request.num_sgs > 0)
+	if (use_bounce_buffer)
+		dma = dep->dwc->bounce_addr;
+	else if (req->request.num_sgs > 0)
 		dma = sg_dma_address(req->start_sg);
 	else
 		dma = req->request.dma;

@@ -1129,56 +1132,35 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 			chain = false;
 
 		if (rem && usb_endpoint_dir_out(dep->endpoint.desc) && !chain) {
-			struct dwc3	*dwc = dep->dwc;
-			struct dwc3_trb	*trb;
-
 			req->needs_extra_trb = true;
 
 			/* prepare normal TRB */
-			dwc3_prepare_one_trb(dep, req, trb_length, true, i);
+			dwc3_prepare_one_trb(dep, req, trb_length,
+					true, i, false);
 
 			/* Now prepare one extra TRB to align transfer size */
-			trb = &dep->trb_pool[dep->trb_enqueue];
-			req->num_trbs++;
-			__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr,
-					maxp - rem, false, 1,
-					req->request.stream_id,
-					req->request.short_not_ok,
-					req->request.no_interrupt,
-					req->request.is_last);
+			dwc3_prepare_one_trb(dep, req, maxp - rem,
+					false, 1, true);
 		} else if (req->request.zero && req->request.length &&
 			   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
 			   !rem && !chain) {
-			struct dwc3	*dwc = dep->dwc;
-			struct dwc3_trb	*trb;
-
 			req->needs_extra_trb = true;
 
 			/* Prepare normal TRB */
-			dwc3_prepare_one_trb(dep, req, trb_length, true, i);
+			dwc3_prepare_one_trb(dep, req, trb_length,
+					true, i, false);
 
 			/* Prepare one extra TRB to handle ZLP */
-			trb = &dep->trb_pool[dep->trb_enqueue];
-			req->num_trbs++;
-			__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr, 0,
-					       !req->direction, 1,
-					       req->request.stream_id,
-					       req->request.short_not_ok,
-					       req->request.no_interrupt,
-					       req->request.is_last);
+			dwc3_prepare_one_trb(dep, req, 0,
+					!req->direction, 1, true);
 
 			/* Prepare one more TRB to handle MPS alignment */
-			if (!req->direction) {
-				trb = &dep->trb_pool[dep->trb_enqueue];
-				req->num_trbs++;
-				__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr, maxp,
-						       false, 1, req->request.stream_id,
-						       req->request.short_not_ok,
-						       req->request.no_interrupt,
-						       req->request.is_last);
-			}
+			if (!req->direction)
+				dwc3_prepare_one_trb(dep, req, maxp,
+						false, 1, true);
 		} else {
-			dwc3_prepare_one_trb(dep, req, trb_length, chain, i);
+			dwc3_prepare_one_trb(dep, req, trb_length,
+					chain, i, false);
 		}
 
 		/*

@@ -1216,54 +1198,29 @@ static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 	unsigned int rem = length % maxp;
 
 	if ((!length || rem) && usb_endpoint_dir_out(dep->endpoint.desc)) {
-		struct dwc3	*dwc = dep->dwc;
-		struct dwc3_trb	*trb;
-
 		req->needs_extra_trb = true;
 
 		/* prepare normal TRB */
-		dwc3_prepare_one_trb(dep, req, length, true, 0);
+		dwc3_prepare_one_trb(dep, req, length, true, 0, false);
 
 		/* Now prepare one extra TRB to align transfer size */
-		trb = &dep->trb_pool[dep->trb_enqueue];
-		req->num_trbs++;
-		__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr, maxp - rem,
-				false, 1, req->request.stream_id,
-				req->request.short_not_ok,
-				req->request.no_interrupt,
-				req->request.is_last);
+		dwc3_prepare_one_trb(dep, req, maxp - rem, false, 1, true);
 	} else if (req->request.zero && req->request.length &&
 		   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
 		   (IS_ALIGNED(req->request.length, maxp))) {
-		struct dwc3	*dwc = dep->dwc;
-		struct dwc3_trb	*trb;
-
 		req->needs_extra_trb = true;
 
 		/* prepare normal TRB */
-		dwc3_prepare_one_trb(dep, req, length, true, 0);
+		dwc3_prepare_one_trb(dep, req, length, true, 0, false);
 
 		/* Prepare one extra TRB to handle ZLP */
-		trb = &dep->trb_pool[dep->trb_enqueue];
-		req->num_trbs++;
-		__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr, 0,
-				!req->direction, 1, req->request.stream_id,
-				req->request.short_not_ok,
-				req->request.no_interrupt,
-				req->request.is_last);
+		dwc3_prepare_one_trb(dep, req, 0, !req->direction, 1, true);
 
 		/* Prepare one more TRB to handle MPS alignment for OUT */
-		if (!req->direction) {
-			trb = &dep->trb_pool[dep->trb_enqueue];
-			req->num_trbs++;
-			__dwc3_prepare_one_trb(dep, trb, dwc->bounce_addr, maxp,
-					       false, 1, req->request.stream_id,
-					       req->request.short_not_ok,
-					       req->request.no_interrupt,
-					       req->request.is_last);
-		}
+		if (!req->direction)
+			dwc3_prepare_one_trb(dep, req, maxp, false, 1, true);
 	} else {
-		dwc3_prepare_one_trb(dep, req, length, false, 0);
+		dwc3_prepare_one_trb(dep, req, length, false, 0, false);
 	}
 }
 

index c0309ad34d28..e49930a7acae 100644
@@ -1132,11 +1132,12 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 			chain = false;
 
 		if (rem && usb_endpoint_dir_out(dep->endpoint.desc) && !chain) {
-			req->needs_extra_trb = true;
-
 			/* prepare normal TRB */
-			dwc3_prepare_one_trb(dep, req, trb_length,
+			if (req->request.length) {
+				req->needs_extra_trb = true;
+				dwc3_prepare_one_trb(dep, req, trb_length,
 					true, i, false);
+			}
 
 			/* Now prepare one extra TRB to align transfer size */
 			dwc3_prepare_one_trb(dep, req, maxp - rem,

@@ -1151,13 +1152,9 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 					true, i, false);
 
 			/* Prepare one extra TRB to handle ZLP */
-			dwc3_prepare_one_trb(dep, req, 0,
-					!req->direction, 1, true);
-
-			/* Prepare one more TRB to handle MPS alignment */
-			if (!req->direction)
-				dwc3_prepare_one_trb(dep, req, maxp,
-						false, 1, true);
+			dwc3_prepare_one_trb(dep, req,
+					req->direction ? 0 : maxp,
+					false, 1, true);
 		} else {
 			dwc3_prepare_one_trb(dep, req, trb_length,
 					chain, i, false);

@@ -1198,10 +1195,11 @@ static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 	unsigned int rem = length % maxp;
 
 	if ((!length || rem) && usb_endpoint_dir_out(dep->endpoint.desc)) {
-		req->needs_extra_trb = true;
-
 		/* prepare normal TRB */
-		dwc3_prepare_one_trb(dep, req, length, true, 0, false);
+		if (req->request.length) {
+			req->needs_extra_trb = true;
+			dwc3_prepare_one_trb(dep, req, length, true, 0, false);
+		}
 
 		/* Now prepare one extra TRB to align transfer size */
 		dwc3_prepare_one_trb(dep, req, maxp - rem, false, 1, true);

@@ -1214,11 +1212,8 @@ static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 		dwc3_prepare_one_trb(dep, req, length, true, 0, false);
 
 		/* Prepare one extra TRB to handle ZLP */
-		dwc3_prepare_one_trb(dep, req, 0, !req->direction, 1, true);
-
-		/* Prepare one more TRB to handle MPS alignment for OUT */
-		if (!req->direction)
-			dwc3_prepare_one_trb(dep, req, maxp, false, 1, true);
+		dwc3_prepare_one_trb(dep, req, req->direction ? 0 : maxp,
+				false, 1, true);
 	} else {
 		dwc3_prepare_one_trb(dep, req, length, false, 0, false);
 	}

@@ -2615,12 +2610,12 @@ static int dwc3_gadget_ep_reclaim_completed_trb(struct dwc3_ep *dep,
 	}
 
 	/*
-	 * If we're dealing with unaligned size OUT transfer, we will be left
-	 * with one TRB pending in the ring. We need to manually clear HWO bit
-	 * from that TRB.
+	 * We use bounce buffer for requests that needs extra TRB or OUT ZLP. If
+	 * this TRB points to the bounce buffer address, it's a MPS alignment
+	 * TRB. Don't add it to req->remaining calculation.
 	 */
-
-	if (req->needs_extra_trb && !(trb->ctrl & DWC3_TRB_CTRL_CHN)) {
+	if (trb->bpl == lower_32_bits(dep->dwc->bounce_addr) &&
+	    trb->bph == upper_32_bits(dep->dwc->bounce_addr)) {
 		trb->ctrl &= ~DWC3_TRB_CTRL_HWO;
 		return 1;
 	}

@@ -2701,17 +2696,8 @@ static int dwc3_gadget_ep_cleanup_completed_request(struct dwc3_ep *dep,
 		goto out;
 
 	if (req->needs_extra_trb) {
-		unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
-
 		ret = dwc3_gadget_ep_reclaim_trb_linear(dep, req, event,
 				status);
-
-		/* Reclaim MPS padding TRB for ZLP */
-		if (!req->direction && req->request.zero && req->request.length &&
-		    !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
-		    (IS_ALIGNED(req->request.length, maxp)))
-			ret = dwc3_gadget_ep_reclaim_trb_linear(dep, req, event, status);
-
 		req->needs_extra_trb = false;
 	}
 

index b0363e3ba4d1..5580caed8b0c 100644
@@ -947,12 +947,16 @@ static void dwc3_ep0_xfer_complete(struct dwc3 *dwc,
 static void __dwc3_ep0_do_control_data(struct dwc3 *dwc,
 		struct dwc3_ep *dep, struct dwc3_request *req)
 {
+	unsigned int		trb_length = 0;
 	int			ret;
 
 	req->direction = !!dep->number;
 
 	if (req->request.length == 0) {
-		dwc3_ep0_prepare_one_trb(dep, dwc->ep0_trb_addr, 0,
+		if (!req->direction)
+			trb_length = dep->endpoint.maxpacket;
+
+		dwc3_ep0_prepare_one_trb(dep, dwc->bounce_addr, trb_length,
 				DWC3_TRBCTL_CONTROL_DATA, false);
 		ret = dwc3_ep0_start_trans(dep);
 	} else if (!IS_ALIGNED(req->request.length, dep->endpoint.maxpacket)

@@ -999,9 +1003,12 @@ static void __dwc3_ep0_do_control_data(struct dwc3 *dwc,
 
 		req->trb = &dwc->ep0_trb[dep->trb_enqueue - 1];
 
+		if (!req->direction)
+			trb_length = dep->endpoint.maxpacket;
+
 		/* Now prepare one extra TRB to align transfer size */
 		dwc3_ep0_prepare_one_trb(dep, dwc->bounce_addr,
-					 0, DWC3_TRBCTL_CONTROL_DATA,
+					 trb_length, DWC3_TRBCTL_CONTROL_DATA,
 					 false);
 		ret = dwc3_ep0_start_trans(dep);
 	} else {

index e49930a7acae..d3896657ddbd 100644
@@ -1094,7 +1094,7 @@ static void dwc3_prepare_one_trb(struct dwc3_ep *dep,
 			stream_id, short_not_ok, no_interrupt, is_last);
 }
 
-static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
+static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {
 	struct scatterlist *sg = req->start_sg;

@@ -1105,6 +1105,7 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 	unsigned int rem = length % maxp;
 	unsigned int remaining = req->request.num_mapped_sgs
 		- req->num_queued_sgs;
+	unsigned int num_trbs = req->num_trbs;
 
 	/*
 	 * If we resume preparing the request, then get the remaining length of

@@ -1131,9 +1132,15 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		if ((i == remaining - 1) || !length)
 			chain = false;
 
+		if (!dwc3_calc_trbs_left(dep))
+			break;
+
 		if (rem && usb_endpoint_dir_out(dep->endpoint.desc) && !chain) {
 			/* prepare normal TRB */
 			if (req->request.length) {
+				if (dwc3_calc_trbs_left(dep) < 2)
+					goto out;
+
 				req->needs_extra_trb = true;
 				dwc3_prepare_one_trb(dep, req, trb_length,
 					true, i, false);

@@ -1145,6 +1152,10 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		} else if (req->request.zero && req->request.length &&
 			   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
 			   !rem && !chain) {
+
+			if (dwc3_calc_trbs_left(dep) < 2)
+				goto out;
+
 			req->needs_extra_trb = true;
 
 			/* Prepare normal TRB */

@@ -1185,18 +1196,28 @@ static void dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		if (!dwc3_calc_trbs_left(dep))
 			break;
 	}
+
+out:
+	return req->num_trbs - num_trbs;
 }
 
-static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
+static int dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {
 	unsigned int length = req->request.length;
 	unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
 	unsigned int rem = length % maxp;
+	unsigned int num_trbs = req->num_trbs;
+
+	if (!dwc3_calc_trbs_left(dep))
+		goto out;
 
 	if ((!length || rem) && usb_endpoint_dir_out(dep->endpoint.desc)) {
 		/* prepare normal TRB */
 		if (req->request.length) {
+			if (dwc3_calc_trbs_left(dep) < 2)
+				goto out;
+
 			req->needs_extra_trb = true;
 			dwc3_prepare_one_trb(dep, req, length, true, 0, false);
 		}

@@ -1206,6 +1227,10 @@ static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 	} else if (req->request.zero && req->request.length &&
 		   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
 		   (IS_ALIGNED(req->request.length, maxp))) {
+
+		if (dwc3_calc_trbs_left(dep) < 2)
+			goto out;
+
 		req->needs_extra_trb = true;
 
 		/* prepare normal TRB */

@@ -1215,8 +1240,14 @@ static void dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 		dwc3_prepare_one_trb(dep, req, req->direction ? 0 : maxp,
 				false, 1, true);
 	} else {
+		if (!dwc3_calc_trbs_left(dep))
+			goto out;
+
 		dwc3_prepare_one_trb(dep, req, length, false, 0, false);
 	}
+
+out:
+	return req->num_trbs - num_trbs;
 }
 
 /*

index d3896657ddbd..f43421246f89 100644
@@ -1257,10 +1257,13 @@ static int dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
  * The function goes through the requests list and sets up TRBs for the
  * transfers. The function returns once there are no more TRBs available or
  * it runs out of requests.
+ *
+ * Returns the number of TRBs prepared or negative errno.
  */
-static void dwc3_prepare_trbs(struct dwc3_ep *dep)
+static int dwc3_prepare_trbs(struct dwc3_ep *dep)
 {
 	struct dwc3_request	*req, *n;
+	int			ret = 0;
 
 	BUILD_BUG_ON_NOT_POWER_OF_2(DWC3_TRB_NUM);
 

@@ -1275,11 +1278,14 @@ static void dwc3_prepare_trbs(struct dwc3_ep *dep)
 	 * break things.
 	 */
 	list_for_each_entry(req, &dep->started_list, list) {
-		if (req->num_pending_sgs > 0)
-			dwc3_prepare_one_trb_sg(dep, req);
+		if (req->num_pending_sgs > 0) {
+			ret = dwc3_prepare_one_trb_sg(dep, req);
+			if (!ret)
+				return ret;
+		}
 
 		if (!dwc3_calc_trbs_left(dep))
-			return;
+			return ret;
 
 		/*
 		 * Don't prepare beyond a transfer. In DWC_usb32, its transfer

@@ -1287,17 +1293,16 @@ static void dwc3_prepare_trbs(struct dwc3_ep *dep)
 		 * active transfer instead of stopping.
 		 */
 		if (dep->stream_capable && req->request.is_last)
-			return;
+			return ret;
 	}
 
 	list_for_each_entry_safe(req, n, &dep->pending_list, list) {
 		struct dwc3	*dwc = dep->dwc;
-		int		ret;
 
 		ret = usb_gadget_map_request_by_dev(dwc->sysdev, &req->request,
 						    dep->direction);
 		if (ret)
-			return;
+			return ret;
 
 		req->sg			= req->request.sg;
 		req->start_sg		= req->sg;

@@ -1305,12 +1310,12 @@ static void dwc3_prepare_trbs(struct dwc3_ep *dep)
 		req->num_pending_sgs	= req->request.num_mapped_sgs;
 
 		if (req->num_pending_sgs > 0)
-			dwc3_prepare_one_trb_sg(dep, req);
+			ret = dwc3_prepare_one_trb_sg(dep, req);
 		else
-			dwc3_prepare_one_trb_linear(dep, req);
+			ret = dwc3_prepare_one_trb_linear(dep, req);
 
-		if (!dwc3_calc_trbs_left(dep))
-			return;
+		if (!ret || !dwc3_calc_trbs_left(dep))
+			return ret;
 
 		/*
 		 * Don't prepare beyond a transfer. In DWC_usb32, its transfer

@@ -1318,8 +1323,10 @@ static void dwc3_prepare_trbs(struct dwc3_ep *dep)
 		 * active transfer instead of stopping.
 		 */
 		if (dep->stream_capable && req->request.is_last)
-			return;
+			return ret;
 	}
+
+	return ret;
 }
 
 static void dwc3_gadget_ep_cleanup_cancelled_requests(struct dwc3_ep *dep);

@@ -1332,12 +1339,12 @@ static int __dwc3_gadget_kick_transfer(struct dwc3_ep *dep)
 	int				ret;
 	u32				cmd;
 
-	if (!dwc3_calc_trbs_left(dep))
-		return 0;
+	ret = dwc3_prepare_trbs(dep);
+	if (ret <= 0)
+		return ret;
 
 	starting = !(dep->flags & DWC3_EP_TRANSFER_STARTED);
 
-	dwc3_prepare_trbs(dep);
 	req = next_request(&dep->started_list);
 	if (!req) {
 		dep->flags |= DWC3_EP_PENDING_REQUEST;

index f43421246f89..840d6e49f668 100644
@@ -1197,7 +1197,27 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 			break;
 	}
 
+	return req->num_trbs - num_trbs;
+
 out:
+	/*
+	 * If we run out of TRBs for MPS alignment setup, then set IOC on the
+	 * previous TRB to get notified for TRB completion to resume when more
+	 * TRBs are available.
+	 *
+	 * Note: normally we shouldn't update the TRB after the HWO bit is set.
+	 * However, the controller doesn't update its internal cache to handle
+	 * the newly prepared TRBs until UPDATE_TRANSFER or START_TRANSFER
+	 * command is executed. At this point, it doesn't happen yet, so we
+	 * should be fine modifying it here.
+	 */
+	if (i) {
+		struct dwc3_trb	*trb;
+
+		trb = dwc3_ep_prev_trb(dep, dep->trb_enqueue);
+		trb->ctrl |= DWC3_TRB_CTRL_IOC;
+	}
+
 	return req->num_trbs - num_trbs;
 }
 

index 840d6e49f668..e6fa2d295439 100644
@@ -1094,6 +1094,47 @@ static void dwc3_prepare_one_trb(struct dwc3_ep *dep,
 			stream_id, short_not_ok, no_interrupt, is_last);
 }
 
+/**
+ * dwc3_prepare_last_sg - prepare TRBs for the last SG entry
+ * @dep: The endpoint that the request belongs to
+ * @req: The request to prepare
+ * @entry_length: The last SG entry size
+ * @node: Indicates whether this is not the first entry (for isoc only)
+ *
+ * Return the number of TRBs prepared.
+ */
+static int dwc3_prepare_last_sg(struct dwc3_ep *dep,
+		struct dwc3_request *req, unsigned int entry_length,
+		unsigned int node)
+{
+	unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
+	unsigned int rem = req->request.length % maxp;
+	unsigned int num_trbs = 1;
+
+	if ((req->request.length && req->request.zero && !rem &&
+			!usb_endpoint_xfer_isoc(dep->endpoint.desc)) ||
+			(!req->direction && rem))
+		num_trbs++;
+
+	if (dwc3_calc_trbs_left(dep) < num_trbs)
+		return 0;
+
+	req->needs_extra_trb = num_trbs > 1;
+
+	/* Prepare a normal TRB */
+	if (req->direction || req->request.length)
+		dwc3_prepare_one_trb(dep, req, entry_length,
+				req->needs_extra_trb, node, false);
+
+	/* Prepare extra TRBs for ZLP and MPS OUT transfer alignment */
+	if ((!req->direction && !req->request.length) || req->needs_extra_trb)
+		dwc3_prepare_one_trb(dep, req,
+				req->direction ? 0 : maxp - rem,
+				false, 1, true);
+
+	return num_trbs;
+}
+
 static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {

@@ -1101,8 +1142,6 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 	struct scatterlist *s;
 	int		i;
 	unsigned int length = req->request.length;
-	unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
-	unsigned int rem = length % maxp;
 	unsigned int remaining = req->request.num_mapped_sgs
 		- req->num_queued_sgs;
 	unsigned int num_trbs = req->num_trbs;

@@ -1116,7 +1155,7 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 
 	for_each_sg(sg, s, remaining, i) {
 		unsigned int trb_length;
-		unsigned int chain = true;
+		bool last_sg = false;
 
 		trb_length = min_t(unsigned int, length, sg_dma_len(s));
 

@@ -1130,45 +1169,16 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		 * mapped sg.
 		 */
 		if ((i == remaining - 1) || !length)
-			chain = false;
+			last_sg = true;
 
 		if (!dwc3_calc_trbs_left(dep))
 			break;
 
-		if (rem && usb_endpoint_dir_out(dep->endpoint.desc) && !chain) {
-			/* prepare normal TRB */
-			if (req->request.length) {
-				if (dwc3_calc_trbs_left(dep) < 2)
-					goto out;
-
-				req->needs_extra_trb = true;
-				dwc3_prepare_one_trb(dep, req, trb_length,
-					true, i, false);
-			}
-
-			/* Now prepare one extra TRB to align transfer size */
-			dwc3_prepare_one_trb(dep, req, maxp - rem,
-					false, 1, true);
-		} else if (req->request.zero && req->request.length &&
-			   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
-			   !rem && !chain) {
-
-			if (dwc3_calc_trbs_left(dep) < 2)
+		if (last_sg) {
+			if (!dwc3_prepare_last_sg(dep, req, trb_length, i))
 				goto out;
-
-			req->needs_extra_trb = true;
-
-			/* Prepare normal TRB */
-			dwc3_prepare_one_trb(dep, req, trb_length,
-					true, i, false);
-
-			/* Prepare one extra TRB to handle ZLP */
-			dwc3_prepare_one_trb(dep, req,
-					req->direction ? 0 : maxp,
-					false, 1, true);
 		} else {
-			dwc3_prepare_one_trb(dep, req, trb_length,
-					chain, i, false);
+			dwc3_prepare_one_trb(dep, req, trb_length, 1, i, false);
 		}
 
 		/*

@@ -1178,7 +1188,7 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 		 * we have free trbs we can continue queuing from where we
 		 * previously stopped
 		 */
-		if (chain)
+		if (!last_sg)
 			req->start_sg = sg_next(s);
 
 		req->num_queued_sgs++;

@@ -1224,50 +1234,7 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 static int dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {
-	unsigned int length = req->request.length;
-	unsigned int maxp = usb_endpoint_maxp(dep->endpoint.desc);
-	unsigned int rem = length % maxp;
-	unsigned int num_trbs = req->num_trbs;
-
-	if (!dwc3_calc_trbs_left(dep))
-		goto out;
-
-	if ((!length || rem) && usb_endpoint_dir_out(dep->endpoint.desc)) {
-		/* prepare normal TRB */
-		if (req->request.length) {
-			if (dwc3_calc_trbs_left(dep) < 2)
-				goto out;
-
-			req->needs_extra_trb = true;
-			dwc3_prepare_one_trb(dep, req, length, true, 0, false);
-		}
-
-		/* Now prepare one extra TRB to align transfer size */
-		dwc3_prepare_one_trb(dep, req, maxp - rem, false, 1, true);
-	} else if (req->request.zero && req->request.length &&
-		   !usb_endpoint_xfer_isoc(dep->endpoint.desc) &&
-		   (IS_ALIGNED(req->request.length, maxp))) {
-
-		if (dwc3_calc_trbs_left(dep) < 2)
-			goto out;
-
-		req->needs_extra_trb = true;
-
-		/* prepare normal TRB */
-		dwc3_prepare_one_trb(dep, req, length, true, 0, false);
-
-		/* Prepare one extra TRB to handle ZLP */
-		dwc3_prepare_one_trb(dep, req, req->direction ? 0 : maxp,
-				false, 1, true);
-	} else {
-		if (!dwc3_calc_trbs_left(dep))
-			goto out;
-
-		dwc3_prepare_one_trb(dep, req, length, false, 0, false);
-	}
-
-out:
-	return req->num_trbs - num_trbs;
+	return dwc3_prepare_last_sg(dep, req, req->request.length, 0);
 }
 
 /*

index e6fa2d295439..ca10619c64a3 100644
@@ -1135,7 +1135,7 @@ static int dwc3_prepare_last_sg(struct dwc3_ep *dep,
 	return num_trbs;
 }
 
-static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
+static int dwc3_prepare_trbs_sg(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {
 	struct scatterlist *sg = req->start_sg;

@@ -1231,7 +1231,7 @@ static int dwc3_prepare_one_trb_sg(struct dwc3_ep *dep,
 	return req->num_trbs - num_trbs;
 }
 
-static int dwc3_prepare_one_trb_linear(struct dwc3_ep *dep,
+static int dwc3_prepare_trbs_linear(struct dwc3_ep *dep,
 		struct dwc3_request *req)
 {
 	return dwc3_prepare_last_sg(dep, req, req->request.length, 0);

@@ -1266,7 +1266,7 @@ static int dwc3_prepare_trbs(struct dwc3_ep *dep)
 	 */
 	list_for_each_entry(req, &dep->started_list, list) {
 		if (req->num_pending_sgs > 0) {
-			ret = dwc3_prepare_one_trb_sg(dep, req);
+			ret = dwc3_prepare_trbs_sg(dep, req);
 			if (!ret)
 				return ret;
 		}

@@ -1297,9 +1297,9 @@ static int dwc3_prepare_trbs(struct dwc3_ep *dep)
 		req->num_pending_sgs	= req->request.num_mapped_sgs;
 
 		if (req->num_pending_sgs > 0)
-			ret = dwc3_prepare_one_trb_sg(dep, req);
+			ret = dwc3_prepare_trbs_sg(dep, req);
 		else
-			ret = dwc3_prepare_one_trb_linear(dep, req);
+			ret = dwc3_prepare_trbs_linear(dep, req);
 
 		if (!ret || !dwc3_calc_trbs_left(dep))
 			return ret;