diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
index d2ea5ce2cefb..d094ec7433f9 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
@@ -3032,12 +3032,7 @@ int amdgpu_vm_init(struct amdgpu_device *adev, struct amdgpu_vm *vm,
 	amdgpu_bo_unreserve(vm->root.base.bo);
 
 	if (pasid) {
-		unsigned long flags;
-
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		r = idr_alloc(&adev->vm_manager.pasid_idr, vm, pasid, pasid + 1,
-			      GFP_ATOMIC);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
+		r = xa_insert_irq(&adev->vm_manager.vms, pasid, vm, GFP_KERNEL);
 		if (r < 0)
 			goto error_free_root;
 
@@ -3047,6 +3042,8 @@ int amdgpu_vm_init(struct amdgpu_device *adev, struct amdgpu_vm *vm,
 	vm->fault_hash = init_fault_hash();
 	if (!vm->fault_hash) {
 		r = -ENOMEM;
+		if (pasid)
+			xa_erase_irq(&adev->vm_manager.vms, pasid);
 		goto error_free_root;
 	}
 
@@ -3104,16 +3101,9 @@ int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm, uns
 	}
 
 	if (pasid) {
-		unsigned long flags;
-
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		r = idr_alloc(&adev->vm_manager.pasid_idr, vm, pasid, pasid + 1,
-			      GFP_ATOMIC);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
-
-		if (r == -ENOSPC)
+		r = xa_insert_irq(&adev->vm_manager.vms, pasid, vm, GFP_KERNEL);
+		if (r < 0)
 			goto unreserve_bo;
-		r = 0;
 	}
 
 	/* Check if PD needs to be reinitialized and do it before
@@ -3124,7 +3114,7 @@ int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm, uns
 			       adev->vm_manager.root_level,
 			       pte_support_ats);
 		if (r)
-			goto free_idr;
+			goto erase;
 	}
 
 	/* Update VM state */
@@ -3137,11 +3127,7 @@ int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm, uns
 		  "CPU update of VM recommended only for large BAR system\n");
 
 	if (vm->pasid) {
-		unsigned long flags;
-
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		idr_remove(&adev->vm_manager.pasid_idr, vm->pasid);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
+		xa_erase_irq(&adev->vm_manager.vms, vm->pasid);
 
 		/* Free the original amdgpu allocated pasid
 		 * Will be replaced with kfd allocated pasid
@@ -3158,14 +3144,9 @@ int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm, uns
 
 	goto unreserve_bo;
 
-free_idr:
-	if (pasid) {
-		unsigned long flags;
-
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		idr_remove(&adev->vm_manager.pasid_idr, pasid);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
-	}
+erase:
+	if (pasid)
+		xa_erase_irq(&adev->vm_manager.vms, pasid);
 unreserve_bo:
 	amdgpu_bo_unreserve(vm->root.base.bo);
 	return r;
@@ -3184,9 +3165,9 @@ void amdgpu_vm_release_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm)
 	if (vm->pasid) {
 		unsigned long flags;
 
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		idr_remove(&adev->vm_manager.pasid_idr, vm->pasid);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
+		xa_lock_irqsave(&adev->vm_manager.vms, flags);
+		__xa_erase(&adev->vm_manager.vms, vm->pasid);
+		xa_unlock_irqrestore(&adev->vm_manager.vms, flags);
 	}
 	vm->pasid = 0;
 }
@@ -3217,9 +3198,9 @@ void amdgpu_vm_fini(struct amdgpu_device *adev, struct amdgpu_vm *vm)
 	if (vm->pasid) {
 		unsigned long flags;
 
-		spin_lock_irqsave(&adev->vm_manager.pasid_lock, flags);
-		idr_remove(&adev->vm_manager.pasid_idr, vm->pasid);
-		spin_unlock_irqrestore(&adev->vm_manager.pasid_lock, flags);
+		xa_lock_irqsave(&adev->vm_manager.vms, flags);
+		__xa_erase(&adev->vm_manager.vms, vm->pasid);
+		xa_unlock_irqrestore(&adev->vm_manager.vms, flags);
 	}
 
 	kfree(vm->fault_hash);
@@ -3299,8 +3280,8 @@ void amdgpu_vm_manager_init(struct amdgpu_device *adev)
 	adev->vm_manager.vm_update_mode = 0;
 #endif
 
-	idr_init(&adev->vm_manager.pasid_idr);
-	spin_lock_init(&adev->vm_manager.pasid_lock);
+	xa_init_flags(&adev->vm_manager.vms,
+			XA_FLAGS_ALLOC | XA_FLAGS_LOCK_IRQ);
 }
 
 /**
@@ -3312,8 +3293,7 @@ void amdgpu_vm_manager_init(struct amdgpu_device *adev)
  */
 void amdgpu_vm_manager_fini(struct amdgpu_device *adev)
 {
-	WARN_ON(!idr_is_empty(&adev->vm_manager.pasid_idr));
-	idr_destroy(&adev->vm_manager.pasid_idr);
+	WARN_ON(!xa_empty(&adev->vm_manager.vms));
 
 	amdgpu_vmid_mgr_fini(adev);
 }
@@ -3364,13 +3344,11 @@ void amdgpu_vm_get_task_info(struct amdgpu_device *adev, unsigned int pasid,
 {
 	struct amdgpu_vm *vm;
 
-	spin_lock(&adev->vm_manager.pasid_lock);
-
-	vm = idr_find(&adev->vm_manager.pasid_idr, pasid);
+	xa_lock(&adev->vm_manager.vms);
+	vm = xa_load(&adev->vm_manager.vms, pasid);
 	if (vm)
 		*task_info = vm->task_info;
-
-	spin_unlock(&adev->vm_manager.pasid_lock);
+	xa_unlock(&adev->vm_manager.vms);
 }
 
 /**
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
index e8dcfd59fc93..84b466eee449 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
@@ -281,8 +281,7 @@ struct amdgpu_vm_manager {
 	/* PASID to VM mapping, will be used in interrupt context to
 	 * look up VM of a page fault
 	 */
-	struct idr				pasid_idr;
-	spinlock_t				pasid_lock;
+	struct xarray				vms;
 };
 
 #define amdgpu_vm_copy_pte(adev, ib, pe, src, count) ((adev)->vm_manager.vm_pte_funcs->copy_pte((ib), (pe), (src), (count)))
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
index bacdaef77b6c..c0e6503b7965 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
@@ -268,11 +268,11 @@ static bool gmc_v9_0_prescreen_iv(struct amdgpu_device *adev,
 		return true;
 
 	/* Track retry faults in per-VM fault FIFO. */
-	spin_lock(&adev->vm_manager.pasid_lock);
-	vm = idr_find(&adev->vm_manager.pasid_idr, entry->pasid);
+	xa_lock(&adev->vm_manager.vms);
+	vm = xa_load(&adev->vm_manager.vms, entry->pasid);
 	if (!vm) {
 		/* VM not found, process it normally */
-		spin_unlock(&adev->vm_manager.pasid_lock);
+		xa_unlock(&adev->vm_manager.vms);
 		return true;
 	}
 
@@ -283,7 +283,7 @@ static bool gmc_v9_0_prescreen_iv(struct amdgpu_device *adev,
 	 * ignore further page faults
 	 */
 	if (r != 0) {
-		spin_unlock(&adev->vm_manager.pasid_lock);
+		xa_unlock(&adev->vm_manager.vms);
 		return false;
 	}
 	/* No locking required with single writer and single reader */
@@ -291,11 +291,11 @@ static bool gmc_v9_0_prescreen_iv(struct amdgpu_device *adev,
 	if (!r) {
 		/* FIFO is full. Ignore it until there is space */
 		amdgpu_vm_clear_fault(vm->fault_hash, key);
-		spin_unlock(&adev->vm_manager.pasid_lock);
+		xa_unlock(&adev->vm_manager.vms);
 		return false;
 	}
 
-	spin_unlock(&adev->vm_manager.pasid_lock);
+	xa_unlock(&adev->vm_manager.vms);
 	/* It's the first fault for this address, process it normally */
 	return true;
 }
