diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index 6d1eeacc559b..5b0a0e9d9c32 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -439,7 +439,7 @@ static int init_shadow_structs(struct kvm *kvm, struct kvm_shadow_vm *vm,
 	return 0;
 }
 
-static bool __exists_shadow(struct kvm *host_kvm)
+static bool exists_shadow(struct kvm *host_kvm)
 {
 	int i;
 	int num_checked = 0;
@@ -463,15 +463,13 @@ static bool __exists_shadow(struct kvm *host_kvm)
  * Return a unique handle to the protected VM on success,
  * negative error code on failure.
  */
-static int insert_shadow_table(struct kvm *kvm, struct kvm_shadow_vm *vm,
-			       size_t shadow_size)
+static int __insert_shadow_table(struct kvm *kvm, struct kvm_shadow_vm *vm,
+			         size_t shadow_size)
 {
 	struct kvm_s2_mmu *mmu = &vm->arch.mmu;
 	int shadow_handle;
 	int vmid;
 
-	hyp_assert_lock_held(&shadow_lock);
-
 	if (unlikely(num_shadow_entries >= KVM_MAX_PVMS))
 		return -ENOMEM;
 
@@ -483,7 +481,7 @@ static int insert_shadow_table(struct kvm *kvm, struct kvm_shadow_vm *vm,
 		return -EINVAL;
 
 	/* Check that a shadow hasn't been created before for this host KVM. */
-	if (unlikely(__exists_shadow(kvm)))
+	if (unlikely(exists_shadow(kvm)))
 		return -EEXIST;
 
 	/* Find the next free entry in the shadow table. */
@@ -511,16 +509,34 @@ static int insert_shadow_table(struct kvm *kvm, struct kvm_shadow_vm *vm,
 	return shadow_handle;
 }
 
+static int insert_shadow_table(struct kvm *kvm, struct kvm_shadow_vm *vm,
+			       size_t shadow_size)
+{
+	int ret;
+
+	hyp_spin_lock(&shadow_lock);
+	ret = __insert_shadow_table(kvm, vm, shadow_size);
+	hyp_spin_unlock(&shadow_lock);
+
+	return ret;
+}
+
 /*
  * Deallocate and remove the shadow table entry corresponding to the handle.
  */
-static void remove_shadow_table(int shadow_handle)
+static void __remove_shadow_table(int shadow_handle)
 {
-	hyp_assert_lock_held(&shadow_lock);
 	shadow_table[shadow_handle_to_index(shadow_handle)] = NULL;
 	num_shadow_entries--;
 }
 
+static void remove_shadow_table(int shadow_handle)
+{
+	hyp_spin_lock(&shadow_lock);
+	__remove_shadow_table(shadow_handle);
+	hyp_spin_unlock(&shadow_lock);
+}
+
 static size_t pkvm_get_shadow_size(int num_vcpus)
 {
 	/* Shadow space for the vm struct and all of its vcpu states. */
@@ -617,29 +633,27 @@ int __pkvm_init_shadow(struct kvm *kvm,
 	if (ret)
 		goto err_remove_pgd;
 
-	ret = init_shadow_structs(kvm, vm, pgd, nr_vcpus);
+	/* Add the entry to the shadow table. */
+	ret = insert_shadow_table(kvm, vm, shadow_size);
 	if (ret < 0)
 		goto err_unpin_host_vcpus;
 
-	/* Add the entry to the shadow table. */
-	hyp_spin_lock(&shadow_lock);
-	ret = insert_shadow_table(kvm, vm, shadow_size);
+	ret = init_shadow_structs(kvm, vm, pgd, nr_vcpus);
 	if (ret < 0)
-		goto err_unlock_unpin_host_vcpus;
+		goto err_remove_shadow_table;
 
 	ret = kvm_guest_prepare_stage2(vm, pgd);
 	if (ret)
 		goto err_remove_shadow_table;
 
-	hyp_spin_unlock(&shadow_lock);
 	return vm->shadow_handle;
 
 err_remove_shadow_table:
 	remove_shadow_table(vm->shadow_handle);
-err_unlock_unpin_host_vcpus:
-	hyp_spin_unlock(&shadow_lock);
+
 err_unpin_host_vcpus:
 	unpin_host_vcpus(vm->shadow_vcpus, nr_vcpus);
+
 err_remove_pgd:
 	WARN_ON(__pkvm_hyp_donate_host(hyp_virt_to_pfn(pgd), nr_pgd_pages));
 
@@ -678,7 +692,7 @@ int __pkvm_teardown_shadow(int shadow_handle)
 		goto err_unlock;
 	}
 
-	remove_shadow_table(shadow_handle);
+	__remove_shadow_table(shadow_handle);
 	hyp_spin_unlock(&shadow_lock);
 
 	/* Reclaim guest pages, and page-table pages */