[PATCH 2/2] drm/amdkfd: Update RDMA implementation such that get_pages() only pins memory of a BO while dma_map() builds an sg_table from the pinned memory

[Ramesh Errabolu]

[Why]
Enable RDMA service when platform has IOMMU ON. The handle of device
to DMA map is available only in dma_map() call context.

[How]
Update implementation get_pages() such that it only Pins memory of a
buffer object. Update implementation dma_map() to use the handle of
DMA device in building an sg_table off of the buffer object

Signed-off-by: Ramesh Errabolu <Ramesh.Errabolu at amd.com>
---
 drivers/gpu/drm/amd/amdkfd/kfd_peerdirect.c | 47 +++++++++++++++++++--
 drivers/gpu/drm/amd/amdkfd/kfd_rdma.c       | 26 +++++++-----
 2 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/drivers/gpu/drm/amd/amdkfd/kfd_peerdirect.c b/drivers/gpu/drm/amd/amdkfd/kfd_peerdirect.c
index 6d7340b101ba..86b4ee710629 100644
--- a/drivers/gpu/drm/amd/amdkfd/kfd_peerdirect.c
+++ b/drivers/gpu/drm/amd/amdkfd/kfd_peerdirect.c
@@ -54,7 +54,6 @@
 #include "kfd_priv.h"
 
 
-
 /* ----------------------- PeerDirect interface ------------------------------*/
 
 /*
@@ -281,6 +280,14 @@ static int amd_get_pages(unsigned long addr, size_t size, int write, int force,
 static int amd_dma_map(struct sg_table *sg_head, void *client_context,
 			struct device *dma_device, int dmasync, int *nmap)
 {
+	struct sg_table *sg_table_tmp;
+	struct kfd_bo *buf_obj;
+	struct kfd_dev *dev;
+	struct kgd_mem *mem;
+	uint64_t offset;
+	uint64_t length;
+	int ret;
+
 	/*
 	 * NOTE/TODO:
 	 * We could have potentially three cases for real memory
@@ -312,15 +319,33 @@ static int amd_dma_map(struct sg_table *sg_head, void *client_context,
 			mem_context->size);
 
 	if (!mem_context->p2p_info) {
-		pr_err("No sg table were allocated\n");
+		pr_err("Buffer was not pinned\n");
 		return -EINVAL;
 	}
 
+	/* Retrieve the handle of buffer object embedded in amd_p2p_info */
+	buf_obj = mem_context->p2p_info->priv;
+	mem = buf_obj->mem;
+	dev = buf_obj->dev;
+	offset = mem_context->va - buf_obj->it.start;
+	length = mem_context->p2p_info->size;
+
+	/* Build sg_table for buffer being exported, including DMA mapping */
+	ret = amdgpu_amdkfd_gpuvm_get_sg_table(dev->kgd, mem,
+			offset, length, &sg_table_tmp);
+	if (ret) {
+		pr_err("Building of sg_table failed\n");
+		return -EFAULT;
+	}
+
+	/* Maintain a copy of the handle to sg_table */
+	mem_context->p2p_info->pages = sg_table_tmp;
+
 	/* Copy information about previosly allocated sg_table */
-	*sg_head = *mem_context->p2p_info->pages;
+	*sg_head = *sg_table_tmp;
 
 	/* Return number of pages */
-	*nmap = mem_context->p2p_info->pages->nents;
+	*nmap = sg_table_tmp->nents;
 
 	return 0;
 }
@@ -328,6 +353,7 @@ static int amd_dma_map(struct sg_table *sg_head, void *client_context,
 static int amd_dma_unmap(struct sg_table *sg_head, void *client_context,
 			   struct device  *dma_device)
 {
+	struct sg_table *sg_table_tmp;
 	struct amd_mem_context *mem_context =
 		(struct amd_mem_context *)client_context;
 
@@ -339,6 +365,19 @@ static int amd_dma_unmap(struct sg_table *sg_head, void *client_context,
 			mem_context->va,
 			mem_context->size);
 
+	/* Determine if sg_table construction was successful */
+	if (!mem_context->p2p_info) {
+		pr_err("Buffer was not pinned\n");
+		return -EINVAL;
+	}
+
+	/* Free sg_table and its elements including DMA unmapping */
+	sg_table_tmp = mem_context->p2p_info->pages;
+	if (sg_table_tmp != NULL) {
+		mem_context->p2p_info->pages = NULL;
+		amdgpu_amdkfd_gpuvm_put_sg_table(sg_table_tmp);
+	}
+
 	/* Assume success */
 	return 0;
 }
diff --git a/drivers/gpu/drm/amd/amdkfd/kfd_rdma.c b/drivers/gpu/drm/amd/amdkfd/kfd_rdma.c
index ec6435aa33bd..67a9213a40b3 100644
--- a/drivers/gpu/drm/amd/amdkfd/kfd_rdma.c
+++ b/drivers/gpu/drm/amd/amdkfd/kfd_rdma.c
@@ -66,7 +66,6 @@ static int get_pages(uint64_t address, uint64_t length, struct pid *pid,
 {
 	struct kfd_bo *buf_obj;
 	struct kgd_mem *mem;
-	struct sg_table *sg_table_tmp;
 	struct kfd_dev *dev;
 	uint64_t last = address + length - 1;
 	uint64_t offset;
@@ -99,10 +98,9 @@ static int get_pages(uint64_t address, uint64_t length, struct pid *pid,
 	dev = buf_obj->dev;
 	offset = address - buf_obj->it.start;
 
-	ret = amdgpu_amdkfd_gpuvm_pin_get_sg_table(dev->kgd, mem,
-			offset, length, &sg_table_tmp);
+	ret = amdgpu_amdkfd_gpuvm_pin_bo(mem);
 	if (ret) {
-		pr_err("amdgpu_amdkfd_gpuvm_pin_get_sg_table failed.\n");
+		pr_err("Pinning of buffer failed.\n");
 		*amd_p2p_data = NULL;
 		goto free_mem;
 	}
@@ -111,7 +109,6 @@ static int get_pages(uint64_t address, uint64_t length, struct pid *pid,
 	rdma_cb_data->amd_p2p_data.size = length;
 	rdma_cb_data->amd_p2p_data.pid = pid;
 	rdma_cb_data->amd_p2p_data.priv = buf_obj;
-	rdma_cb_data->amd_p2p_data.pages = sg_table_tmp;
 	rdma_cb_data->amd_p2p_data.kfd_proc = p;
 
 	rdma_cb_data->free_callback = free_callback;
@@ -140,24 +137,33 @@ static int put_pages_helper(struct amd_p2p_info *p2p_data)
 {
 	struct kfd_bo *buf_obj;
 	struct kfd_dev *dev;
-	struct sg_table *sg_table_tmp;
 	struct rdma_cb *rdma_cb_data;
+	struct sg_table *sg_table_tmp;
 
 	if (!p2p_data) {
 		pr_err("amd_p2p_info pointer is invalid.\n");
 		return -EINVAL;
 	}
 
-	rdma_cb_data = container_of(p2p_data, struct rdma_cb, amd_p2p_data);
+	/* Determine if sg_table been released */
+	sg_table_tmp = p2p_data->pages;
+	if (sg_table_tmp != NULL) {
+		p2p_data->pages = NULL;
+		amdgpu_amdkfd_gpuvm_put_sg_table(sg_table_tmp);
+	}
 
+	/* Unpin pages of buffer object */
 	buf_obj = p2p_data->priv;
-	dev = buf_obj->dev;
-	sg_table_tmp = p2p_data->pages;
+	p2p_data->priv = NULL;
+	if (buf_obj != NULL)
+		amdgpu_amdkfd_gpuvm_unpin_bo(buf_obj->mem);
 
+	 /* Remove callback info from list of callback objects */
+	rdma_cb_data = container_of(p2p_data, struct rdma_cb, amd_p2p_data);
 	list_del(&rdma_cb_data->node);
 	kfree(rdma_cb_data);
 
-	amdgpu_amdkfd_gpuvm_unpin_put_sg_table(buf_obj->mem, sg_table_tmp);
+	dev = buf_obj->dev;
 	kfd_dec_compute_active(dev);
 
 	return 0;
-- 
2.29.2