diff --git a/drivers/staging/android/ion/ion_dma_buf.c b/drivers/staging/android/ion/ion_dma_buf.c
index c3734ff40940..5a2e17917099 100644
--- a/drivers/staging/android/ion/ion_dma_buf.c
+++ b/drivers/staging/android/ion/ion_dma_buf.c
@@ -287,11 +287,16 @@ static void *ion_dma_buf_vmap(struct dma_buf *dmabuf)
 {
 	struct ion_buffer *buffer = dmabuf->priv;
 	struct ion_heap *heap = buffer->heap;
+	void *vaddr;
 
-	if (!heap->buf_ops.vmap)
-		return ERR_PTR(-EOPNOTSUPP);
+	if (heap->buf_ops.vmap)
+		return heap->buf_ops.vmap(dmabuf);
 
-	return heap->buf_ops.vmap(dmabuf);
+	mutex_lock(&buffer->lock);
+	vaddr = ion_buffer_kmap_get(buffer);
+	mutex_unlock(&buffer->lock);
+
+	return vaddr;
 }
 
 static void ion_dma_buf_vunmap(struct dma_buf *dmabuf, void *vaddr)
@@ -299,10 +304,14 @@ static void ion_dma_buf_vunmap(struct dma_buf *dmabuf, void *vaddr)
 	struct ion_buffer *buffer = dmabuf->priv;
 	struct ion_heap *heap = buffer->heap;
 
-	if (!heap->buf_ops.vunmap)
+	if (heap->buf_ops.vunmap) {
+		heap->buf_ops.vunmap(dmabuf, vaddr);
 		return;
+	}
 
-	return heap->buf_ops.vunmap(dmabuf, vaddr);
+	mutex_lock(&buffer->lock);
+	ion_buffer_kmap_put(buffer);
+	mutex_unlock(&buffer->lock);
 }
 
 static int ion_dma_buf_get_flags(struct dma_buf *dmabuf, unsigned long *flags)