[PATCH 4/4] io_uring/zcrx: coalesce areas with huge pages

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

Try to shrink the page array into fewer larger folios if possible. This
reduces the footprint for the array and prepares us for future huge page
optimisations.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
---
 io_uring/zcrx.c | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/io_uring/zcrx.c b/io_uring/zcrx.c
index 784c4ed6c780..fd0d97830854 100644
--- a/io_uring/zcrx.c
+++ b/io_uring/zcrx.c
@@ -219,6 +219,8 @@ static int io_import_area_memory(struct io_zcrx_ifq *ifq,
 				 struct io_zcrx_area *area,
 				 struct io_uring_zcrx_area_reg *area_reg)
 {
+	struct io_imu_folio_data data;
+	bool coalesced = false;
 	struct iovec iov;
 	int nr_pages;
 	int ret;
@@ -239,6 +241,21 @@ static int io_import_area_memory(struct io_zcrx_ifq *ifq,
 	area->nr_folios = nr_pages;
 	area->folio_shift = PAGE_SHIFT;
 	area->chunk_id_offset = 0;
+
+	if (nr_pages > 1 && io_check_coalesce_buffer(area->pages, nr_pages, &data)) {
+		if (data.nr_pages_mid != 1)
+			coalesced = io_coalesce_buffer(&area->pages, &nr_pages, &data);
+	}
+
+	if (coalesced) {
+		size_t folio_size = 1UL << data.folio_shift;
+		size_t offset = folio_size - (data.nr_pages_head << PAGE_SHIFT);
+
+		area->nr_folios = nr_pages;
+		area->folio_shift = data.folio_shift;
+		area->chunk_id_offset = offset >> PAGE_SHIFT;
+	}
+
 	return 0;
 }
 
-- 
2.48.1