[PATCH 06/13] PCI/P2PDMA: add align parameter for pci_p2pdma_add_resource()

[Hou Tao <houtao@huaweicloud.com>]

From: Hou Tao <houtao1@huawei.com>

The align parameter is used to align both the mapping of p2p dma memory
and to enable the compound page for p2p dma memory in the kernel and in
the userspace.

Signed-off-by: Hou Tao <houtao1@huawei.com>
---
 drivers/accel/habanalabs/common/hldio.c |  3 +-
 drivers/nvme/host/pci.c                 |  2 +-
 drivers/pci/p2pdma.c                    | 38 +++++++++++++++++++++----
 include/linux/pci-p2pdma.h              |  4 +--
 4 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/drivers/accel/habanalabs/common/hldio.c b/drivers/accel/habanalabs/common/hldio.c
index 083ae5610875..4d1528dbde9f 100644
--- a/drivers/accel/habanalabs/common/hldio.c
+++ b/drivers/accel/habanalabs/common/hldio.c
@@ -372,7 +372,8 @@ int hl_p2p_region_init(struct hl_device *hdev, struct hl_p2p_region *p2pr)
 	int rc, i;
 
 	/* Start by publishing our p2p memory */
-	rc = pci_p2pdma_add_resource(hdev->pdev, p2pr->bar, p2pr->size, p2pr->bar_offset);
+	rc = pci_p2pdma_add_resource(hdev->pdev, p2pr->bar, p2pr->size, PAGE_SIZE,
+				     p2pr->bar_offset);
 	if (rc) {
 		dev_err(hdev->dev, "error adding p2p resource: %d\n", rc);
 		goto err;
diff --git a/drivers/nvme/host/pci.c b/drivers/nvme/host/pci.c
index 0e4caeab739c..b070095bae5e 100644
--- a/drivers/nvme/host/pci.c
+++ b/drivers/nvme/host/pci.c
@@ -2309,7 +2309,7 @@ static void nvme_map_cmb(struct nvme_dev *dev)
 			     dev->bar + NVME_REG_CMBMSC);
 	}
 
-	if (pci_p2pdma_add_resource(pdev, bar, size, offset)) {
+	if (pci_p2pdma_add_resource(pdev, bar, size, PAGE_SIZE, offset)) {
 		dev_warn(dev->ctrl.device,
 			 "failed to register the CMB\n");
 		hi_lo_writeq(0, dev->bar + NVME_REG_CMBMSC);
diff --git a/drivers/pci/p2pdma.c b/drivers/pci/p2pdma.c
index dd64ec830fdd..70482e240304 100644
--- a/drivers/pci/p2pdma.c
+++ b/drivers/pci/p2pdma.c
@@ -23,6 +23,7 @@
 
 struct pci_p2pdma {
 	struct gen_pool *pool;
+	size_t align;
 	bool p2pmem_published;
 	struct xarray map_types;
 	struct p2pdma_provider mem[PCI_STD_NUM_BARS];
@@ -211,7 +212,7 @@ static void p2pdma_folio_free(struct folio *folio)
 	struct percpu_ref *ref;
 
 	gen_pool_free_owner(p2pdma->pool, (uintptr_t)page_to_virt(page),
-			    PAGE_SIZE, (void **)&ref);
+			    p2pdma->align, (void **)&ref);
 	percpu_ref_put(ref);
 }
 
@@ -323,17 +324,22 @@ struct p2pdma_provider *pcim_p2pdma_provider(struct pci_dev *pdev, int bar)
 }
 EXPORT_SYMBOL_GPL(pcim_p2pdma_provider);
 
-static int pci_p2pdma_setup_pool(struct pci_dev *pdev)
+static int pci_p2pdma_setup_pool(struct pci_dev *pdev, size_t align)
 {
 	struct pci_p2pdma *p2pdma;
 	int ret;
 
 	p2pdma = rcu_dereference_protected(pdev->p2pdma, 1);
-	if (p2pdma->pool)
+	if (p2pdma->pool) {
+		/* Two p2pdma BARs with different alignment ? */
+		if (p2pdma->align != align)
+			return -EINVAL;
 		/* We already setup pools, do nothing, */
 		return 0;
+	}
 
-	p2pdma->pool = gen_pool_create(PAGE_SHIFT, dev_to_node(&pdev->dev));
+	p2pdma->align = align;
+	p2pdma->pool = gen_pool_create(ilog2(p2pdma->align), dev_to_node(&pdev->dev));
 	if (!p2pdma->pool)
 		return -ENOMEM;
 
@@ -363,18 +369,31 @@ static void pci_p2pdma_unmap_mappings(void *data)
 				     p2pmem_group.name);
 }
 
+static inline int pci_p2pdma_check_pagemap_align(struct pci_dev *pdev, int bar,
+						 u64 size, size_t align,
+						 u64 offset)
+{
+	if (align == PAGE_SIZE)
+		return 0;
+	return -EINVAL;
+}
+
 /**
  * pci_p2pdma_add_resource - add memory for use as p2p memory
  * @pdev: the device to add the memory to
  * @bar: PCI BAR to add
  * @size: size of the memory to add, may be zero to use the whole BAR
+ * @align: dev memory mapping alignment of the memory to add. It is used
+ * to optimize the mappings both in userspace and kernel space when
+ * transparent huge page is supported. The possible values are
+ * PAGE_SIZE, PMD_SIZE, and PUD_SIZE.
  * @offset: offset into the PCI BAR
  *
  * The memory will be given ZONE_DEVICE struct pages so that it may
  * be used with any DMA request.
  */
 int pci_p2pdma_add_resource(struct pci_dev *pdev, int bar, size_t size,
-			    u64 offset)
+			    size_t align, u64 offset)
 {
 	struct pci_p2pdma_pagemap *p2p_pgmap;
 	struct p2pdma_provider *mem;
@@ -395,11 +414,18 @@ int pci_p2pdma_add_resource(struct pci_dev *pdev, int bar, size_t size,
 	if (size + offset > pci_resource_len(pdev, bar))
 		return -EINVAL;
 
+	error = pci_p2pdma_check_pagemap_align(pdev, bar, size, align, offset);
+	if (error) {
+		pci_info_ratelimited(pdev, "invalid align 0x%zx for bar %d\n",
+				     align, bar);
+		return error;
+	}
+
 	error = pcim_p2pdma_init(pdev);
 	if (error)
 		return error;
 
-	error = pci_p2pdma_setup_pool(pdev);
+	error = pci_p2pdma_setup_pool(pdev, align);
 	if (error)
 		return error;
 
diff --git a/include/linux/pci-p2pdma.h b/include/linux/pci-p2pdma.h
index 517e121d2598..2fa671274c45 100644
--- a/include/linux/pci-p2pdma.h
+++ b/include/linux/pci-p2pdma.h
@@ -69,7 +69,7 @@ enum pci_p2pdma_map_type {
 int pcim_p2pdma_init(struct pci_dev *pdev);
 struct p2pdma_provider *pcim_p2pdma_provider(struct pci_dev *pdev, int bar);
 int pci_p2pdma_add_resource(struct pci_dev *pdev, int bar, size_t size,
-		u64 offset);
+			    size_t align, u64 offset);
 int pci_p2pdma_distance_many(struct pci_dev *provider, struct device **clients,
 			     int num_clients, bool verbose);
 struct pci_dev *pci_p2pmem_find_many(struct device **clients, int num_clients);
@@ -97,7 +97,7 @@ static inline struct p2pdma_provider *pcim_p2pdma_provider(struct pci_dev *pdev,
 	return NULL;
 }
 static inline int pci_p2pdma_add_resource(struct pci_dev *pdev, int bar,
-		size_t size, u64 offset)
+		size_t size, size_t align, u64 offset)
 {
 	return -EOPNOTSUPP;
 }
-- 
2.29.2