1 files changed, 34 insertions, 10 deletions

diff --git a/remoteproc/ethosu_remoteproc.c b/remoteproc/ethosu_remoteproc.c
index c6d7735..344627d 100644
--- a/remoteproc/ethosu_remoteproc.c
+++ b/remoteproc/ethosu_remoteproc.c
@@ -122,6 +122,38 @@ static void ethosu_mbox_top(struct mbox_client *client,
 	queue_work(erproc->wq, &erproc->work);
 }
 
+static int ethosu_mem_alloc(struct rproc *rproc,
+			    struct rproc_mem_entry *mem)
+{
+	struct device *dev = rproc->dev.parent;
+	void *va;
+
+	va = (__force void *)devm_ioremap_wc(dev, mem->dma, mem->len);
+	if (IS_ERR_OR_NULL(va)) {
+		dev_err(dev, "Failed to remap address. pa=%pa, len=%zu",
+			&mem->dma,
+			mem->len);
+
+		return -ENOMEM;
+	}
+
+	mem->va = va;
+	mem->is_iomem = true;
+
+	return 0;
+}
+
+static int ethosu_mem_release(struct rproc *rproc,
+			      struct rproc_mem_entry *mem)
+{
+	struct device *dev = rproc->dev.parent;
+
+	if (mem->va)
+		devm_iounmap(dev, (__force __iomem void *)mem->va);
+
+	return 0;
+}
+
 static int ethosu_add_carveout(struct rproc *rproc,
 			       const phys_addr_t pa,
 			       const size_t size,
@@ -129,7 +161,6 @@ static int ethosu_add_carveout(struct rproc *rproc,
 {
 	struct device *dev = rproc->dev.parent;
 	dma_addr_t da;
-	void __iomem *va;
 	struct rproc_mem_entry *mem;
 
 	da = translate_phys_to_dma(dev, pa);
@@ -141,15 +172,8 @@ static int ethosu_add_carveout(struct rproc *rproc,
 		return -ENOMEM;
 	}
 
-	va = devm_ioremap_wc(dev, pa, size);
-	if (!va) {
-		dev_err(dev, "Failed to remap address. pa=%pa, len=%zu", &pa,
-			size);
-
-		return -ENOMEM;
-	}
-
-	mem = rproc_mem_entry_init(dev, va, pa, size, da, NULL, NULL, name);
+	mem = rproc_mem_entry_init(dev, NULL, pa, size, da, ethosu_mem_alloc,
+				   ethosu_mem_release, name);
 	if (!mem)
 		return -ENOMEM;