[PATCH v4 2/6] KVM: arm64: nv: Avoid full shadow s2 unmap

[Oliver Upton]

Hey,

On Tue, Jul 14, 2026 at 12:59:21PM +0100, Wei-Lin Chang wrote:
> +static bool valid_entry(unsigned long entry)
> +{
> +	WARN_ON(entry & VALID_ENTRY && entry & UNKNOWN_IPA);
> +	return entry & VALID_ENTRY;
> +}
> +
> +static bool unknown_ipa_entry(unsigned long entry)
> +{
> +	WARN_ON(entry & VALID_ENTRY && entry & UNKNOWN_IPA);
> +	return entry & UNKNOWN_IPA;
> +}

These should be WARN_ON_ONCE(), but I find the condition you're
asserting to be a bit confusing. An aliased reverse map entry is still a
valid entry, would you not set the valid bit?

> +void kvm_record_nested_revmap(gpa_t canonical_ipa, struct kvm_s2_mmu *mmu,
> +			      gpa_t nested_ipa, size_t map_size)
> +{
> +	struct maple_tree *revmap_mt = &mmu->nested_revmap_mt;
> +	gpa_t canonical_ipa_end;
> +	u64 entry, new_entry = 0;
> +
> +	lockdep_assert_held_read(&kvm_s2_mmu_to_kvm(mmu)->mmu_lock);
> +
> +	if (mmu->nested_revmap_broken)
> +		return;
> +
> +	if (WARN_ON(!IS_ALIGNED(canonical_ipa, map_size)))
> +		canonical_ipa = ALIGN_DOWN(canonical_ipa, map_size);
> +
> +	canonical_ipa_end = canonical_ipa + map_size - 1;
> +	MA_STATE(mas_rmap, revmap_mt, canonical_ipa, canonical_ipa_end);
> +
> +	mtree_lock(revmap_mt);
> +	entry = xa_to_value(mas_find(&mas_rmap, canonical_ipa_end));
> +
> +	if (entry) {
> +		/* parallel faults can be adding the same mapping */
> +		if (valid_entry(entry) &&
> +		    mas_rmap.index == canonical_ipa &&
> +		    mas_rmap.last == canonical_ipa_end &&
> +		    nested_ipa == (entry & ADDR_MASK))
> +			goto unlock;
> +		/*
> +		 * Create a "UNKNOWN_IPA" range that spans all the overlapping
> +		 * ranges and store it.
> +		 */
> +		while (entry && mas_rmap.index <= canonical_ipa_end) {
> +			canonical_ipa = min(mas_rmap.index, canonical_ipa);
> +			canonical_ipa_end = max(mas_rmap.last, canonical_ipa_end);
> +			entry = xa_to_value(mas_find(&mas_rmap, canonical_ipa_end));
> +		}
> +		new_entry |= UNKNOWN_IPA;
> +	} else {
> +		new_entry |= nested_ipa;
> +		new_entry |= VALID_ENTRY;
> +	}
> +
> +	mas_set_range(&mas_rmap, canonical_ipa, canonical_ipa_end);
> +	if (mas_store_gfp(&mas_rmap, xa_mk_value(new_entry),
> +			  GFP_NOWAIT | __GFP_ACCOUNT))

The general pattern for handling this situation would be to preallocate
nodes prior to acquiring the lock. If we can preallocate then we can
just outright refuse to create a shadow stage-2 mapping if the update to
the reverse map fails.

> +		mmu->nested_revmap_broken = true;
> +unlock:
> +	mtree_unlock(revmap_mt);
> +}
> +
>  void kvm_init_nested_s2_mmu(struct kvm_s2_mmu *mmu)
>  {
>  	/* CnP being set denotes an invalid entry */
>  	mmu->tlb_vttbr = VTTBR_CNP_BIT;
>  	mmu->nested_stage2_enabled = false;
>  	atomic_set(&mmu->refcnt, 0);
> +	mt_init(&mmu->nested_revmap_mt);
> +	mmu->nested_revmap_broken = false;
>  }
>  
>  void kvm_vcpu_load_hw_mmu(struct kvm_vcpu *vcpu)
> @@ -1225,6 +1305,89 @@ void kvm_nested_s2_wp(struct kvm *kvm)
>  	kvm_invalidate_vncr_ipa(kvm, 0, BIT(kvm->arch.mmu.pgt->ia_bits));
>  }
>  
> +static void reset_revmap_and_unmap(struct kvm_s2_mmu *mmu, bool may_block)
> +{
> +	mtree_destroy(&mmu->nested_revmap_mt);
> +	mmu->nested_revmap_broken = false;
> +	kvm_stage2_unmap_range(mmu, 0, kvm_phys_size(mmu), may_block);
> +}
> +
> +static void unmap_mmu_cipa_range(struct kvm_s2_mmu *mmu, gpa_t canonical_ipa,
> +				  size_t unmap_size, bool may_block)
> +{
> +	struct maple_tree *revmap_mt = &mmu->nested_revmap_mt;
> +	gpa_t canonical_ipa_end = canonical_ipa + unmap_size - 1;
> +	size_t entry_size;
> +	gpa_t next_addr;
> +	u64 entry;
> +	MA_STATE(mas_rmap, revmap_mt, canonical_ipa, canonical_ipa_end);
> +
> +	lockdep_assert_held_write(&kvm_s2_mmu_to_kvm(mmu)->mmu_lock);
> +
> +	if (mmu->nested_revmap_broken) {
> +		reset_revmap_and_unmap(mmu, may_block);
> +		return;
> +	}
> +
> +	if (!mmu->nested_stage2_enabled) {
> +		kvm_stage2_unmap_range(mmu, canonical_ipa, unmap_size, may_block);
> +		return;
> +	}
> +
> +	mtree_lock(revmap_mt);
> +	entry = xa_to_value(mas_find(&mas_rmap, canonical_ipa_end));
> +
> +	while (entry && mas_rmap.index <= canonical_ipa_end) {
> +		entry_size = mas_rmap.last - mas_rmap.index + 1;
> +		next_addr = mas_rmap.index + entry_size;
> +		/*
> +		 * Give up and invalidate this s2 mmu if the unmap range
> +		 * touches any UNKNOWN_IPA range.
> +		 */
> +		if (unknown_ipa_entry(entry)) {
> +			mtree_unlock(revmap_mt);
> +			reset_revmap_and_unmap(mmu, may_block);
> +			return;
> +		}
> +
> +		/*
> +		 * Ignore result, it is okay if a reverse mapping erase
> +		 * fails.
> +		 */
> +		mas_store_gfp(&mas_rmap, NULL, GFP_NOWAIT | __GFP_ACCOUNT);

I wonder if we can handle this a bit more gracefully.

Ideally we'd be able to place a search mark on the entry flagging it as
stale which could be used to avoid attempting to overinvalidate the
shadow stage-2.

Interestingly enough, Liam recently gave the suggestion [*] of storing
XA_ZERO_ENTRY for a similar use case where erasure happens in an atomic
context. We should be able to do the exact same thing here and push the
cleanup to the next insertion.

[*] https://lore.kernel.org/all/iv7sxqezzz42xieytbf3lskutricz7ltkr3wii3w4jbnbzto7f@2tz5hwa43l37/

Thanks,
Oliver