1 files changed, 23 insertions, 12 deletions

diff --git a/drivers/infiniband/hw/hns/hns_roce_alloc.c b/drivers/infiniband/hw/hns/hns_roce_alloc.c
index e1b433cdd5e2..3e4c5253ab5c 100644
--- a/drivers/infiniband/hw/hns/hns_roce_alloc.c
+++ b/drivers/infiniband/hw/hns/hns_roce_alloc.c
@@ -67,6 +67,7 @@ void hns_roce_bitmap_free(struct hns_roce_bitmap *bitmap, unsigned long obj,
 {
 	hns_roce_bitmap_free_range(bitmap, obj, 1, rr);
 }
+EXPORT_SYMBOL_GPL(hns_roce_bitmap_free);
 
 int hns_roce_bitmap_alloc_range(struct hns_roce_bitmap *bitmap, int cnt,
 				int align, unsigned long *obj)
@@ -160,39 +161,47 @@ void hns_roce_buf_free(struct hns_roce_dev *hr_dev, u32 size,
 		       struct hns_roce_buf *buf)
 {
 	int i;
-	struct device *dev = &hr_dev->pdev->dev;
+	struct device *dev = hr_dev->dev;
 	u32 bits_per_long = BITS_PER_LONG;
 
 	if (buf->nbufs == 1) {
 		dma_free_coherent(dev, size, buf->direct.buf, buf->direct.map);
 	} else {
-		if (bits_per_long == 64)
+		if (bits_per_long == 64 && buf->page_shift == PAGE_SHIFT)
 			vunmap(buf->direct.buf);
 
 		for (i = 0; i < buf->nbufs; ++i)
 			if (buf->page_list[i].buf)
-				dma_free_coherent(&hr_dev->pdev->dev, PAGE_SIZE,
+				dma_free_coherent(dev, 1 << buf->page_shift,
 						  buf->page_list[i].buf,
 						  buf->page_list[i].map);
 		kfree(buf->page_list);
 	}
 }
+EXPORT_SYMBOL_GPL(hns_roce_buf_free);
 
 int hns_roce_buf_alloc(struct hns_roce_dev *hr_dev, u32 size, u32 max_direct,
-		       struct hns_roce_buf *buf)
+		       struct hns_roce_buf *buf, u32 page_shift)
 {
 	int i = 0;
 	dma_addr_t t;
 	struct page **pages;
-	struct device *dev = &hr_dev->pdev->dev;
+	struct device *dev = hr_dev->dev;
 	u32 bits_per_long = BITS_PER_LONG;
+	u32 page_size = 1 << page_shift;
+	u32 order;
 
 	/* SQ/RQ buf lease than one page, SQ + RQ = 8K */
 	if (size <= max_direct) {
 		buf->nbufs = 1;
 		/* Npages calculated by page_size */
-		buf->npages = 1 << get_order(size);
-		buf->page_shift = PAGE_SHIFT;
+		order = get_order(size);
+		if (order <= page_shift - PAGE_SHIFT)
+			order = 0;
+		else
+			order -= page_shift - PAGE_SHIFT;
+		buf->npages = 1 << order;
+		buf->page_shift = page_shift;
 		/* MTT PA must be recorded in 4k alignment, t is 4k aligned */
 		buf->direct.buf = dma_alloc_coherent(dev, size, &t, GFP_KERNEL);
 		if (!buf->direct.buf)
@@ -207,9 +216,9 @@ int hns_roce_buf_alloc(struct hns_roce_dev *hr_dev, u32 size, u32 max_direct,
 
 		memset(buf->direct.buf, 0, size);
 	} else {
-		buf->nbufs = (size + PAGE_SIZE - 1) / PAGE_SIZE;
+		buf->nbufs = (size + page_size - 1) / page_size;
 		buf->npages = buf->nbufs;
-		buf->page_shift = PAGE_SHIFT;
+		buf->page_shift = page_shift;
 		buf->page_list = kcalloc(buf->nbufs, sizeof(*buf->page_list),
 					 GFP_KERNEL);
 
@@ -218,16 +227,16 @@ int hns_roce_buf_alloc(struct hns_roce_dev *hr_dev, u32 size, u32 max_direct,
 
 		for (i = 0; i < buf->nbufs; ++i) {
 			buf->page_list[i].buf = dma_alloc_coherent(dev,
-								  PAGE_SIZE, &t,
+								  page_size, &t,
 								  GFP_KERNEL);
 
 			if (!buf->page_list[i].buf)
 				goto err_free;
 
 			buf->page_list[i].map = t;
-			memset(buf->page_list[i].buf, 0, PAGE_SIZE);
+			memset(buf->page_list[i].buf, 0, page_size);
 		}
-		if (bits_per_long == 64) {
+		if (bits_per_long == 64 && page_shift == PAGE_SHIFT) {
 			pages = kmalloc_array(buf->nbufs, sizeof(*pages),
 					      GFP_KERNEL);
 			if (!pages)
@@ -241,6 +250,8 @@ int hns_roce_buf_alloc(struct hns_roce_dev *hr_dev, u32 size, u32 max_direct,
 			kfree(pages);
 			if (!buf->direct.buf)
 				goto err_free;
+		} else {
+			buf->direct.buf = NULL;
 		}
 	}