In preparation for using the TLB tags allocator for SVM, which has a
range of ASIDs allocated for SEV/SNP, pass in the number of reserved TLB
tags when initializing the TLB tags allocator. The bitmap is
conceptually shifted such that bit=0 corresponds to tag=nr_reserved.

Specifying the minimum value during initialization also makes the API
clearer, as tag=0 is now explicitly reserved by the caller.

No functional change intended for VMX as VPID=0 was already implicitly
reserved.

Signed-off-by: Yosry Ahmed <yosry@kernel.org>
---
 arch/x86/kvm/mmu.h     |  2 +-
 arch/x86/kvm/mmu/mmu.c | 42 ++++++++++++++++++++++++++++--------------
 arch/x86/kvm/vmx/vmx.h |  2 +-
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/arch/x86/kvm/mmu.h b/arch/x86/kvm/mmu.h
index de79e002edf8f..b443a5083bfb0 100644
--- a/arch/x86/kvm/mmu.h
+++ b/arch/x86/kvm/mmu.h
@@ -413,7 +413,7 @@ static inline bool kvm_is_gfn_alias(struct kvm *kvm, gfn_t gfn)
 
 typedef unsigned int kvm_tlb_tag_t;
 
-int kvm_init_tlb_tags(unsigned int nr);
+int kvm_init_tlb_tags(unsigned int nr_total, unsigned int nr_reserved);
 void kvm_destroy_tlb_tags(void);
 kvm_tlb_tag_t kvm_alloc_tlb_tag(void);
 void kvm_free_tlb_tag(kvm_tlb_tag_t tag);
diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index d9edba502cac7..05d2ba2598e1b 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -8073,9 +8073,17 @@ static struct {
 	spinlock_t	lock;
 	unsigned long	*bitmap;
 	unsigned int	nr;
+	kvm_tlb_tag_t	start;
 } tlb_tags;
 
-int kvm_init_tlb_tags(unsigned int nr)
+/*
+ * Initialize @nr_total TLB tags, where @nr_reserved of those tags are reserved.
+ * Reserved tags start at 0 and are contiguous. Hence, usable tags are
+ * [nr_reserved, nr_total). At least one reserved tag is required, as 0 is the
+ * host TLB tag on both VMX and AMD, and is the value returned on a failed tag
+ * allocation.
+ */
+int kvm_init_tlb_tags(unsigned int nr_total, unsigned int nr_reserved)
 {
 	/*
 	 * Limit the number of TLB tags to VMX's hardcoded maximum of 0x10000
@@ -8085,7 +8093,12 @@ int kvm_init_tlb_tags(unsigned int nr)
 	 * need a solution that works for both Intel and AMD.
 	 */
 	const unsigned int MAX_NR_TLB_TAGS = VMX_NR_VPIDS;
+	unsigned int nr;
 
+	if (WARN_ON_ONCE(!nr_reserved || nr_reserved > nr_total))
+		return -EINVAL;
+
+	nr = nr_total - nr_reserved;
 	if (!nr)
 		return 0;
 
@@ -8099,13 +8112,8 @@ int kvm_init_tlb_tags(unsigned int nr)
 	if (!tlb_tags.bitmap)
 		return -ENOMEM;
 
-	/*
-	 * 0 is the host's TLB tag for both VMX's VPID and SVM's ASID, and is
-	 * returned on failed allocations (e.g. no more tags left).
-	 */
-	__set_bit(0, tlb_tags.bitmap);
-
 	tlb_tags.nr = nr;
+	tlb_tags.start = nr_reserved;
 	spin_lock_init(&tlb_tags.lock);
 	return 0;
 }
@@ -8120,30 +8128,36 @@ EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_destroy_tlb_tags);
 
 kvm_tlb_tag_t kvm_alloc_tlb_tag(void)
 {
-	kvm_tlb_tag_t tag;
+	unsigned int bit;
 
 	if (!tlb_tags.bitmap)
 		return 0;
 
 	guard(spinlock)(&tlb_tags.lock);
 
-	tag = find_first_zero_bit(tlb_tags.bitmap, tlb_tags.nr);
-	if (tag >= tlb_tags.nr)
+	bit = find_first_zero_bit(tlb_tags.bitmap, tlb_tags.nr);
+	if (bit >= tlb_tags.nr)
 		return 0;
 
-	__set_bit(tag, tlb_tags.bitmap);
-	return tag;
+	__set_bit(bit, tlb_tags.bitmap);
+	return tlb_tags.start + bit;
 }
 EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_alloc_tlb_tag);
 
 void kvm_free_tlb_tag(kvm_tlb_tag_t tag)
 {
-	if (!tag || WARN_ON_ONCE(tag >= tlb_tags.nr))
+	unsigned int bit;
+
+	if (!tag || WARN_ON_ONCE(tag < tlb_tags.start))
+		return;
+
+	bit = tag - tlb_tags.start;
+	if (WARN_ON_ONCE(bit >= tlb_tags.nr))
 		return;
 
 	guard(spinlock)(&tlb_tags.lock);
 
-	__clear_bit(tag, tlb_tags.bitmap);
+	__clear_bit(bit, tlb_tags.bitmap);
 }
 EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_free_tlb_tag);
 
diff --git a/arch/x86/kvm/vmx/vmx.h b/arch/x86/kvm/vmx/vmx.h
index 3de3ee53ccbdc..7b7e57fd549c4 100644
--- a/arch/x86/kvm/vmx/vmx.h
+++ b/arch/x86/kvm/vmx/vmx.h
@@ -343,7 +343,7 @@ static __always_inline u32 vmx_get_intr_info(struct kvm_vcpu *vcpu)
 
 static __always_inline int init_vpids(void)
 {
-	return enable_vpid ? kvm_init_tlb_tags(VMX_NR_VPIDS) : 0;
+	return enable_vpid ? kvm_init_tlb_tags(VMX_NR_VPIDS, 1) : 0;
 }
 
 static __always_inline void destroy_vpids(void)
-- 
2.55.0.229.g6434b31f56-goog