diff --git a/kernel/ethosu_buffer.c b/kernel/ethosu_buffer.c
index bf7d745..69b5007 100644
--- a/kernel/ethosu_buffer.c
+++ b/kernel/ethosu_buffer.c
@@ -24,6 +24,7 @@
 #include "ethosu_buffer.h"
 
 #include "ethosu_device.h"
+#include "ethosu_dma_mem.h"
 #include "uapi/ethosu.h"
 
 #include <linux/anon_inodes.h>
@@ -71,9 +72,7 @@
 
 	dev_dbg(dev, "Buffer destroy. buf=0x%pK", buf);
 
-	memset(buf->cpu_addr, 0, buf->size);
-	dma_free_coherent(dev, buf->size, buf->cpu_addr,
-			  buf->dma_addr);
+	ethosu_dma_mem_free(&buf->dma_mem);
 
 	memset(buf, 0, sizeof(*buf));
 	devm_kfree(dev, buf);
@@ -103,8 +102,8 @@
 	dev_dbg(dev, "Buffer mmap. file=0x%pK, buf=0x%pK\n",
 		file, buf);
 
-	ret = dma_mmap_coherent(dev, vma, buf->cpu_addr,
-				buf->dma_addr, buf->size);
+	ret = dma_mmap_coherent(dev, vma, buf->dma_mem->cpu_addr,
+				buf->dma_mem->dma_addr, buf->dma_mem->size);
 
 	return ret;
 }
@@ -126,7 +125,7 @@
 	 */
 	switch (whence) {
 	case SEEK_END:
-		return buf->size;
+		return buf->dma_mem->size;
 	case SEEK_SET:
 		return 0;
 	default:
@@ -148,13 +147,13 @@
 		return -ENOMEM;
 
 	buf->dev = dev;
-	buf->size = size;
 	kref_init(&buf->kref);
 
-	buf->cpu_addr = dma_alloc_coherent(dev, size, &buf->dma_addr,
-					   GFP_KERNEL);
-	if (!buf->cpu_addr)
+	buf->dma_mem = ethosu_dma_mem_alloc(dev, size);
+	if (IS_ERR(buf->dma_mem)) {
+		ret = PTR_ERR(buf->dma_mem);
 		goto free_buf;
+	}
 
 	ret = anon_inode_getfd("ethosu-buffer", &ethosu_buffer_fops, buf,
 			       O_RDWR | O_CLOEXEC);
@@ -168,14 +167,13 @@
 
 	dev_dbg(dev,
 		"Buffer create. file=0x%pK, fd=%d, buf=0x%pK, size=%zu, cpu_addr=0x%pK, dma_addr=0x%llx, phys_addr=0x%llx\n",
-		buf->file, ret, buf, size, buf->cpu_addr, buf->dma_addr,
-		virt_to_phys(buf->cpu_addr));
+		buf->file, ret, buf, size, buf->dma_mem->cpu_addr,
+		buf->dma_mem->dma_addr, virt_to_phys(buf->dma_mem->cpu_addr));
 
 	return ret;
 
 free_dma:
-	dma_free_coherent(dev, buf->size, buf->cpu_addr,
-			  buf->dma_addr);
+	ethosu_dma_mem_free(&buf->dma_mem);
 
 free_buf:
 	memset(buf, 0, sizeof(*buf));