diff --git a/gxp-common-platform.c b/gxp-common-platform.c
index 7514b11..6318e98 100644
--- a/gxp-common-platform.c
+++ b/gxp-common-platform.c
@@ -280,8 +280,7 @@
 	return ret;
 }
 
-static int gxp_unmap_buffer(struct gxp_client *client,
-			    struct gxp_map_ioctl __user *argp)
+static int gxp_unmap_buffer(struct gxp_client *client, struct gxp_map_ioctl __user *argp)
 {
 	struct gxp_dev *gxp = client->gxp;
 	struct gxp_map_ioctl ibuf;
@@ -300,29 +299,32 @@
 		goto out;
 	}
 
-	map = gxp_vd_mapping_search(client->vd,
-				    (dma_addr_t)ibuf.device_address);
+	down_write(&client->vd->mappings_semaphore);
+
+	map = gxp_vd_mapping_search_locked(client->vd, (dma_addr_t)ibuf.device_address);
 	if (!map) {
-		dev_err(gxp->dev,
-			"Mapping not found for provided device address %#llX\n",
+		dev_err(gxp->dev, "Mapping not found for provided device address %#llX\n",
 			ibuf.device_address);
 		ret = -EINVAL;
-		goto out;
 	} else if (!map->host_address) {
 		dev_err(gxp->dev, "dma-bufs must be unmapped via GXP_UNMAP_DMABUF\n");
 		ret = -EINVAL;
-		goto out;
 	}
 
-	WARN_ON(map->host_address != ibuf.host_address);
+	if (ret) {
+		up_write(&client->vd->mappings_semaphore);
+		goto out_put;
+	}
 
-	gxp_vd_mapping_remove(client->vd, map);
-	gxp_mapping_iova_log(client, map,
-			     GXP_IOVA_LOG_UNMAP | GXP_IOVA_LOG_BUFFER);
+	gxp_vd_mapping_remove_locked(client->vd, map);
+	up_write(&client->vd->mappings_semaphore);
 
+	gxp_mapping_iova_log(client, map, GXP_IOVA_LOG_UNMAP | GXP_IOVA_LOG_BUFFER);
+
+out_put:
 	/* Release the reference from gxp_vd_mapping_search() */
-	gxp_mapping_put(map);
-
+	if (map)
+		gxp_mapping_put(map);
 out:
 	up_read(&client->semaphore);
 
@@ -1335,8 +1337,7 @@
 	return ret;
 }
 
-static int gxp_unmap_dmabuf(struct gxp_client *client,
-			    struct gxp_map_dmabuf_ioctl __user *argp)
+static int gxp_unmap_dmabuf(struct gxp_client *client, struct gxp_map_dmabuf_ioctl __user *argp)
 {
 	struct gxp_dev *gxp = client->gxp;
 	struct gxp_map_dmabuf_ioctl ibuf;
@@ -1355,14 +1356,17 @@
 		goto out;
 	}
 
+	down_write(&client->vd->mappings_semaphore);
+
 	/*
 	 * Fetch and remove the internal mapping records.
 	 * If host_address is not 0, the provided device_address belongs to a
 	 * non-dma-buf mapping.
 	 */
-	mapping = gxp_vd_mapping_search(client->vd, ibuf.device_address);
+	mapping = gxp_vd_mapping_search_locked(client->vd, ibuf.device_address);
 	if (IS_ERR_OR_NULL(mapping) || mapping->host_address) {
 		dev_warn(gxp->dev, "No dma-buf mapped for given IOVA\n");
+		up_write(&client->vd->mappings_semaphore);
 		/*
 		 * If the device address belongs to a non-dma-buf mapping,
 		 * release the reference to it obtained via the search.
@@ -1374,10 +1378,10 @@
 	}
 
 	/* Remove the mapping from its VD, releasing the VD's reference */
-	gxp_vd_mapping_remove(client->vd, mapping);
+	gxp_vd_mapping_remove_locked(client->vd, mapping);
+	up_write(&client->vd->mappings_semaphore);
 
-	gxp_mapping_iova_log(client, mapping,
-			     GXP_IOVA_LOG_UNMAP | GXP_IOVA_LOG_DMABUF);
+	gxp_mapping_iova_log(client, mapping, GXP_IOVA_LOG_UNMAP | GXP_IOVA_LOG_DMABUF);
 
 	/* Release the reference from gxp_vd_mapping_search() */
 	gxp_mapping_put(mapping);
diff --git a/gxp-vd.c b/gxp-vd.c
index 37d8ab5..a0bfe95 100644
--- a/gxp-vd.c
+++ b/gxp-vd.c
@@ -1315,14 +1315,19 @@
 			   struct gxp_mapping *map)
 {
 	down_write(&vd->mappings_semaphore);
+	gxp_vd_mapping_remove_locked(vd, map);
+	up_write(&vd->mappings_semaphore);
+}
+
+void gxp_vd_mapping_remove_locked(struct gxp_virtual_device *vd, struct gxp_mapping *map)
+{
+	lockdep_assert_held_write(&vd->mappings_semaphore);
 
 	/* Drop the mapping from this virtual device's records */
 	rb_erase(&map->node, &vd->mappings_root);
 
 	/* Release the reference obtained in gxp_vd_mapping_store() */
 	gxp_mapping_put(map);
-
-	up_write(&vd->mappings_semaphore);
 }
 
 static bool is_device_address_in_mapping(struct gxp_mapping *mapping,
@@ -1339,7 +1344,7 @@
 	struct rb_node *node;
 	struct gxp_mapping *mapping;
 
-	down_read(&vd->mappings_semaphore);
+	lockdep_assert_held(&vd->mappings_semaphore);
 
 	node = vd->mappings_root.rb_node;
 
@@ -1349,7 +1354,6 @@
 		    (check_range &&
 		     is_device_address_in_mapping(mapping, device_address))) {
 			gxp_mapping_get(mapping);
-			up_read(&vd->mappings_semaphore);
 			return mapping; /* Found it */
 		} else if (mapping->device_address > device_address) {
 			node = node->rb_left;
@@ -1358,14 +1362,24 @@
 		}
 	}
 
-	up_read(&vd->mappings_semaphore);
-
 	return NULL;
 }
 
 struct gxp_mapping *gxp_vd_mapping_search(struct gxp_virtual_device *vd,
 					  dma_addr_t device_address)
 {
+	struct gxp_mapping *mapping;
+
+	down_read(&vd->mappings_semaphore);
+	mapping = gxp_vd_mapping_search_locked(vd, device_address);
+	up_read(&vd->mappings_semaphore);
+
+	return mapping;
+}
+
+struct gxp_mapping *gxp_vd_mapping_search_locked(struct gxp_virtual_device *vd,
+						 dma_addr_t device_address)
+{
 	return gxp_vd_mapping_internal_search(vd, device_address, false);
 }
 
@@ -1373,7 +1387,13 @@
 gxp_vd_mapping_search_in_range(struct gxp_virtual_device *vd,
 			       dma_addr_t device_address)
 {
-	return gxp_vd_mapping_internal_search(vd, device_address, true);
+	struct gxp_mapping *mapping;
+
+	down_read(&vd->mappings_semaphore);
+	mapping = gxp_vd_mapping_internal_search(vd, device_address, true);
+	up_read(&vd->mappings_semaphore);
+
+	return mapping;
 }
 
 struct gxp_mapping *gxp_vd_mapping_search_host(struct gxp_virtual_device *vd,
diff --git a/gxp-vd.h b/gxp-vd.h
index e0bd5ce..47ba60a 100644
--- a/gxp-vd.h
+++ b/gxp-vd.h
@@ -265,6 +265,12 @@
 			   struct gxp_mapping *map);
 
 /**
+ * gxp_vd_mapping_remove_locked() - The same as `gxp_vd_mapping_remove` but the caller holds
+ *                                  @vd->mappings_semaphore as write.
+ */
+void gxp_vd_mapping_remove_locked(struct gxp_virtual_device *vd, struct gxp_mapping *map);
+
+/**
  * gxp_vd_mapping_search() - Obtain a reference to the mapping starting at the
  *                           specified device address
  * @vd: The virtual device to search for the mapping
@@ -278,6 +284,13 @@
 					  dma_addr_t device_address);
 
 /**
+ * gxp_vd_mapping_search_locked() - The same as `gxp_vd_mapping_search` but the caller holds
+ *                                  @vd->mappings_semaphore.
+ */
+struct gxp_mapping *gxp_vd_mapping_search_locked(struct gxp_virtual_device *vd,
+						 dma_addr_t device_address);
+
+/**
  * gxp_vd_mapping_search_in_range() - Obtain a reference to the mapping which
  *                                    contains the specified device address
  * @vd: The virtual device to search for the mapping