[PATCH 6/8] io_uring/zcrx: add infra for large pages

[Pavel Begunkov <asml.silence@gmail.com>]

Currently, zcrx assumes PAGE_SIZE pages, add infrastructure to support
uniformly sized higher order pages. sg_table hides most of the details
and the only place that need index recalculation is the copy fallback.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
---
 io_uring/zcrx.c | 13 ++++++++-----
 io_uring/zcrx.h |  1 +
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/io_uring/zcrx.c b/io_uring/zcrx.c
index a83b80c16491..44b5f8084279 100644
--- a/io_uring/zcrx.c
+++ b/io_uring/zcrx.c
@@ -44,13 +44,14 @@ static inline struct folio *io_niov_folio(const struct net_iov *niov,
 					  unsigned long *off)
 {
 	struct io_zcrx_area *area = io_zcrx_iov_to_area(niov);
-	struct page *page;
+	unsigned long niov_off, folio_idx;
 
 	lockdep_assert(!area->mem.is_dmabuf);
 
-	page = area->mem.pages[net_iov_idx(niov)];
-	*off = (page - compound_head(page)) << PAGE_SHIFT;
-	return page_folio(page);
+	niov_off = net_iov_idx(niov) << PAGE_SHIFT;
+	folio_idx = niov_off >> area->mem.folio_shift;
+	*off = niov_off - (folio_idx << area->mem.folio_shift);
+	return page_folio(area->mem.pages[folio_idx]);
 }
 
 static int io_populate_area_dma(struct io_zcrx_ifq *ifq,
@@ -164,6 +165,7 @@ static int io_import_umem(struct io_zcrx_ifq *ifq,
 			  struct io_zcrx_mem *mem,
 			  struct io_uring_zcrx_area_reg *area_reg)
 {
+	unsigned folio_shift = PAGE_SHIFT;
 	struct scatterlist *sgl;
 	struct page **pages;
 	int nr_pages, ret;
@@ -182,11 +184,12 @@ static int io_import_umem(struct io_zcrx_ifq *ifq,
 	if (ret)
 		return ret;
 	for_each_sg(mem->page_sg_table.sgl, sgl, nr_pages, i)
-		sg_set_page(sgl, pages[i], PAGE_SIZE, 0);
+		sg_set_page(sgl, pages[i], 1U << folio_shift, 0);
 
 	mem->pages = pages;
 	mem->nr_folios = nr_pages;
 	mem->size = area_reg->len;
+	mem->folio_shift = folio_shift;
 	return 0;
 }
 
diff --git a/io_uring/zcrx.h b/io_uring/zcrx.h
index 89015b923911..4f718b3088d9 100644
--- a/io_uring/zcrx.h
+++ b/io_uring/zcrx.h
@@ -14,6 +14,7 @@ struct io_zcrx_mem {
 
 	struct page			**pages;
 	unsigned long			nr_folios;
+	unsigned			folio_shift;
 	struct sg_table			page_sg_table;
 
 	struct dma_buf_attachment	*attach;
-- 
2.49.0