diff --git a/drivers/dma-buf/heaps/cma_heap.c b/drivers/dma-buf/heaps/cma_heap.c
index 4772c8729e77..00921db42bf7 100644
--- a/drivers/dma-buf/heaps/cma_heap.c
+++ b/drivers/dma-buf/heaps/cma_heap.c
@@ -199,7 +199,7 @@ static void *cma_heap_do_vmap(struct cma_heap_buffer *buffer)
 	return vaddr;
 }
 
-static void *cma_heap_vmap(struct dma_buf *dmabuf)
+static int cma_heap_vmap(struct dma_buf *dmabuf, struct dma_buf_map *map)
 {
 	struct cma_heap_buffer *buffer = dmabuf->priv;
 	void *vaddr;
@@ -212,18 +212,21 @@ static void *cma_heap_vmap(struct dma_buf *dmabuf)
 	}
 
 	vaddr = cma_heap_do_vmap(buffer);
-	if (IS_ERR(vaddr))
-		goto out;
+	if (IS_ERR(vaddr)) {
+		mutex_unlock(&buffer->lock);
+		return PTR_ERR(vaddr);
+	}
 
 	buffer->vaddr = vaddr;
 	buffer->vmap_cnt++;
 out:
 	mutex_unlock(&buffer->lock);
 
-	return vaddr;
+	dma_buf_map_set_vaddr(map, vaddr);
+	return 0;
 }
 
-static void cma_heap_vunmap(struct dma_buf *dmabuf, void *vaddr)
+static void cma_heap_vunmap(struct dma_buf *dmabuf, struct dma_buf_map *map)
 {
 	struct cma_heap_buffer *buffer = dmabuf->priv;