1 files changed, 9 insertions, 51 deletions

diff --git a/drivers/misc/habanalabs/gaudi/gaudi.c b/drivers/misc/habanalabs/gaudi/gaudi.c
index fda3d8a85ada..49d4b5dda115 100644
--- a/drivers/misc/habanalabs/gaudi/gaudi.c
+++ b/drivers/misc/habanalabs/gaudi/gaudi.c
@@ -7755,9 +7755,6 @@ static int gaudi_internal_cb_pool_init(struct hl_device *hdev,
 		struct hl_ctx *ctx)
 {
 	struct gaudi_device *gaudi = hdev->asic_specific;
-	bool flush_pte;
-	u64 va, pa;
-	s64 off;
 	int min_alloc_order, rc, collective_cb_size;
 
 	if (!(gaudi->hw_cap_initialized & HW_CAP_MMU))
@@ -7802,48 +7799,23 @@ static int gaudi_internal_cb_pool_init(struct hl_device *hdev,
 		goto destroy_internal_cb_pool;
 
 	mutex_lock(&ctx->mmu_lock);
-
-	/* The mapping is done page by page since we can't assure allocated ptr
-	 * is aligned to HOST_SPACE_INTERNAL_CB_SZ
-	 */
-	for (off = 0 ; off < HOST_SPACE_INTERNAL_CB_SZ ; off += PAGE_SIZE_4KB) {
-		va = hdev->internal_cb_va_base + off;
-		pa = hdev->internal_cb_pool_dma_addr + off;
-		flush_pte = (off + PAGE_SIZE_4KB) >= HOST_SPACE_INTERNAL_CB_SZ;
-		rc = hl_mmu_map(ctx, va, pa, PAGE_SIZE_4KB, flush_pte);
-		if (rc) {
-			dev_err(hdev->dev,
-				"Map failed for va 0x%llx to pa 0x%llx\n",
-				va, pa);
-			goto unmap;
-		}
-	}
+	rc = hl_mmu_map_contiguous(ctx, hdev->internal_cb_va_base,
+			hdev->internal_cb_pool_dma_addr,
+			HOST_SPACE_INTERNAL_CB_SZ);
 
 	hdev->asic_funcs->mmu_invalidate_cache(hdev, false, VM_TYPE_USERPTR);
-
 	mutex_unlock(&ctx->mmu_lock);
 
-	return 0;
+	if (rc)
+		goto unreserve_internal_cb_pool;
 
-unmap:
-	for (; off >= 0 ; off -= PAGE_SIZE_4KB) {
-		va = hdev->internal_cb_va_base + off;
-		flush_pte = (off - (s32) PAGE_SIZE_4KB) < 0;
-		if (hl_mmu_unmap(ctx, va, PAGE_SIZE_4KB, flush_pte))
-			dev_warn_ratelimited(hdev->dev,
-					"failed to unmap va 0x%llx\n", va);
-	}
+	return 0;
 
+unreserve_internal_cb_pool:
 	hl_unreserve_va_block(hdev, ctx, hdev->internal_cb_va_base,
 			HOST_SPACE_INTERNAL_CB_SZ);
-
-	hdev->asic_funcs->mmu_invalidate_cache(hdev, true, VM_TYPE_USERPTR);
-
-	mutex_unlock(&ctx->mmu_lock);
-
 destroy_internal_cb_pool:
 	gen_pool_destroy(hdev->internal_cb_pool);
-
 free_internal_cb_pool:
 	hdev->asic_funcs->asic_dma_free_coherent(hdev,
 			HOST_SPACE_INTERNAL_CB_SZ,
@@ -7857,30 +7829,16 @@ static void gaudi_internal_cb_pool_fini(struct hl_device *hdev,
 		struct hl_ctx *ctx)
 {
 	struct gaudi_device *gaudi = hdev->asic_specific;
-	bool flush_pte = false;
-	u64 va, off;
 
 	if (!(gaudi->hw_cap_initialized & HW_CAP_MMU))
 		return;
 
 	mutex_lock(&ctx->mmu_lock);
-
-	for (off = 0 ; off < HOST_SPACE_INTERNAL_CB_SZ ; off += PAGE_SIZE_4KB) {
-		va = hdev->internal_cb_va_base + off;
-
-		if (off + PAGE_SIZE_4KB >= HOST_SPACE_INTERNAL_CB_SZ)
-			flush_pte = true;
-
-		if (hl_mmu_unmap(ctx, va, PAGE_SIZE_4KB, flush_pte))
-			dev_warn_ratelimited(hdev->dev,
-					"failed to unmap va 0x%llx\n", va);
-	}
-
+	hl_mmu_unmap_contiguous(ctx, hdev->internal_cb_va_base,
+			HOST_SPACE_INTERNAL_CB_SZ);
 	hl_unreserve_va_block(hdev, ctx, hdev->internal_cb_va_base,
 			HOST_SPACE_INTERNAL_CB_SZ);
-
 	hdev->asic_funcs->mmu_invalidate_cache(hdev, true, VM_TYPE_USERPTR);
-
 	mutex_unlock(&ctx->mmu_lock);
 
 	gen_pool_destroy(hdev->internal_cb_pool);