diff --git a/kernel/ethosu_buffer.c b/kernel/ethosu_buffer.c
index ac32aca..a83a95a 100644
--- a/kernel/ethosu_buffer.c
+++ b/kernel/ethosu_buffer.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022 Arm Limited.
+ * Copyright 2020-2023 Arm Limited and/or its affiliates
  *
  * This program is free software and is provided to you under the terms of the
  * GNU General Public License version 2 as published by the Free Software
@@ -32,6 +32,7 @@
 #include <linux/of_address.h>
 #include <linux/file.h>
 #include <linux/fs.h>
+#include <linux/remoteproc.h>
 #include <linux/uaccess.h>
 
 /****************************************************************************
@@ -61,59 +62,22 @@
  * Functions
  ****************************************************************************/
 
-/*
- * The 'dma-ranges' device tree property for shared dma memory does not seem
- * to be fully supported for coherent memory. Therefor we apply the DMA range
- * offset ourselves.
- */
-static dma_addr_t ethosu_buffer_dma_ranges(struct device *dev,
-					   dma_addr_t dma_addr,
-					   size_t dma_buf_size)
+__attribute__((used))
+static dma_addr_t ethosu_pa_to_da(struct device *dev,
+				  phys_addr_t pa,
+				  size_t len)
 {
-	struct device_node *node = dev->of_node;
-	const __be32 *ranges;
-	int len;
-	int naddr;
-	int nsize;
-	int inc;
-	int i;
+	struct rproc *rproc = rproc_get_by_child(dev);
+	struct rproc_mem_entry *mem;
 
-	if (!node)
-		return dma_addr;
+	list_for_each_entry(mem, &rproc->carveouts, node) {
+		ssize_t offset = pa - mem->dma;
 
-	/* Get the #address-cells and #size-cells properties */
-	naddr = of_n_addr_cells(node);
-	nsize = of_n_size_cells(node);
-
-	/* Read the 'dma-ranges' property */
-	ranges = of_get_property(node, "dma-ranges", &len);
-	if (!ranges || len <= 0)
-		return dma_addr;
-
-	dev_dbg(dev, "ranges=%p, len=%d, naddr=%d, nsize=%d\n",
-		ranges, len, naddr, nsize);
-
-	len /= sizeof(*ranges);
-	inc = naddr + naddr + nsize;
-
-	for (i = 0; (i + inc) <= len; i += inc) {
-		dma_addr_t daddr;
-		dma_addr_t paddr;
-		dma_addr_t size;
-
-		daddr = of_read_number(&ranges[i], naddr);
-		paddr = of_read_number(&ranges[i + naddr], naddr);
-		size = of_read_number(&ranges[i + naddr + naddr], nsize);
-
-		dev_dbg(dev, "daddr=0x%llx, paddr=0x%llx, size=0x%llx\n",
-			daddr, paddr, size);
-
-		if (dma_addr >= paddr &&
-		    (dma_addr + dma_buf_size) < (paddr + size))
-			return dma_addr + daddr - paddr;
+		if (offset >= 0 && offset + len <= mem->len)
+			return mem->da + offset;
 	}
 
-	return dma_addr;
+	return (dma_addr_t)-1;
 }
 
 static bool ethosu_buffer_verify(struct file *file)
@@ -128,8 +92,8 @@
 
 	dev_info(buf->edev->dev, "Buffer destroy. buf=0x%pK\n", buf);
 
-	dma_free_coherent(buf->edev->dev, buf->capacity, buf->cpu_addr,
-			  buf->dma_addr_orig);
+	dma_free_coherent(buf->dev, buf->capacity, buf->cpu_addr,
+			  buf->dma_addr);
 
 	devm_kfree(buf->edev->dev, buf);
 }
@@ -157,8 +121,7 @@
 		 file, buf);
 
 	ret = dma_mmap_coherent(buf->edev->dev, vma, buf->cpu_addr,
-				buf->dma_addr_orig,
-				buf->capacity);
+				buf->dma_addr, buf->capacity);
 
 	return ret;
 }
@@ -224,6 +187,8 @@
 int ethosu_buffer_create(struct ethosu_device *edev,
 			 size_t capacity)
 {
+	struct rproc *rproc = rproc_get_by_child(edev->dev);
+	struct device *dev = rproc->dev.parent;
 	struct ethosu_buffer *buf;
 	int ret = -ENOMEM;
 
@@ -235,20 +200,17 @@
 		return -ENOMEM;
 
 	buf->edev = edev;
+	buf->dev = dev;
 	buf->capacity = capacity;
 	buf->offset = 0;
 	buf->size = 0;
 	kref_init(&buf->kref);
 
-	buf->cpu_addr = dma_alloc_coherent(buf->edev->dev, capacity,
-					   &buf->dma_addr_orig, GFP_KERNEL);
+	buf->cpu_addr = dma_alloc_coherent(dev, capacity,
+					   &buf->dma_addr, GFP_KERNEL);
 	if (!buf->cpu_addr)
 		goto free_buf;
 
-	buf->dma_addr = ethosu_buffer_dma_ranges(buf->edev->dev,
-						 buf->dma_addr_orig,
-						 buf->capacity);
-
 	ret = anon_inode_getfd("ethosu-buffer", &ethosu_buffer_fops, buf,
 			       O_RDWR | O_CLOEXEC);
 	if (ret < 0)
@@ -258,15 +220,15 @@
 	fput(buf->file);
 
 	dev_info(buf->edev->dev,
-		 "Buffer create. file=0x%pK, fd=%d, buf=0x%pK, capacity=%zu, cpu_addr=0x%pK, dma_addr=0x%llx, dma_addr_orig=0x%llx, phys_addr=0x%llx\n",
+		 "Buffer create. file=0x%pK, fd=%d, buf=0x%pK, capacity=%zu, cpu_addr=0x%pK, dma_addr=0x%llx, phys_addr=0x%llx\n",
 		 buf->file, ret, buf, capacity, buf->cpu_addr, buf->dma_addr,
-		 buf->dma_addr_orig, virt_to_phys(buf->cpu_addr));
+		 virt_to_phys(buf->cpu_addr));
 
 	return ret;
 
 free_dma:
 	dma_free_coherent(buf->edev->dev, buf->capacity, buf->cpu_addr,
-			  buf->dma_addr_orig);
+			  buf->dma_addr);
 
 free_buf:
 	devm_kfree(buf->edev->dev, buf);