diff --git a/drivers/iommu/pkvm-pviommu.c b/drivers/iommu/pkvm-pviommu.c
index 5b793b84ea74..15073c25c6be 100644
--- a/drivers/iommu/pkvm-pviommu.c
+++ b/drivers/iommu/pkvm-pviommu.c
@@ -92,21 +92,24 @@ static void pviommu_domain_remove_map(struct pviommu_domain *pv_domain,
 	/* Range can cover multiple entries. */
 	while (start < end) {
 		MA_STATE(mas, &pv_domain->mappings, start, end);
-		u64 entry = xa_to_value(mas_find(&mas, start));
+		u64 entry;
 		u64 old_start, old_end;
 
+		mtree_lock(mas.tree);
+		entry = xa_to_value(mas_find(&mas, start));
 		old_start = mas.index;
 		old_end = mas.last;
 		mas_erase(&mas);
-		/* Insert the rest if not removed. */
-		if (start > old_start)
-			mtree_store_range(&pv_domain->mappings, old_start, start - 1,
-					  xa_mk_value(entry), GFP_KERNEL);
-
-		if (old_end > end)
-			mtree_store_range(&pv_domain->mappings, end + 1, old_end,
-					  xa_mk_value(entry + end - old_start + 1), GFP_KERNEL);
-
+		if (start > old_start) {
+			MA_STATE(mas_border, &pv_domain->mappings, old_start, start - 1);
+			WARN_ON(mas_store_gfp(&mas_border, xa_mk_value(entry), GFP_ATOMIC));
+		}
+		if (old_end > end) {
+			MA_STATE(mas_border, &pv_domain->mappings, end + 1, old_end);
+			WARN_ON(mas_store_gfp(&mas_border, xa_mk_value(entry + end - old_start + 1),
+				GFP_ATOMIC));
+		}
+		mtree_unlock(mas.tree);
 		start = old_end + 1;
 	}
 }
@@ -114,8 +117,11 @@ static void pviommu_domain_remove_map(struct pviommu_domain *pv_domain,
 static u64 pviommu_domain_find(struct pviommu_domain *pv_domain, u64 key)
 {
 	MA_STATE(mas, &pv_domain->mappings, key, key);
-	void *entry = mas_find(&mas, key);
+	void *entry;
 
+	mtree_lock(mas.tree);
+	entry = mas_find(&mas, key);
+	mtree_unlock(mas.tree);
 	/* No entry. */
 	if (!xa_is_value(entry))
 		return 0;