[PATCH v4 01/17] x86/virt/tdx: Enhance tdx_pamt_get/put() to support huge pages

[Yan Zhao <yan.y.zhao@intel.com>]

The TDX module uses PAMT to track metadata for every physical page (within
TDMR) across three levels: 1GB, 2MB, and 4KB. The VMM is responsible for
providing backing pages for the PAMT. Without Dynamic PAMT (DPAMT), all
PAMT backing pages are allocated statically at system boot time. With DPAMT
enabled, PAMT backing pages tracking physical pages at the 4KB level
(referred to as DPAMT pages) are dynamically installed and uninstalled at
runtime.

Currently, the VMM invokes tdx_pamt_get()/tdx_pamt_put() to install and
uninstall DPAMT pages under the assumption that all physical pages are
tracked at 4KB level.  However, when a guest page is mapped as a huge page
in the S-EPT, the TDX module tracks it directly at the huge page level,
meaning installation and uninstallation of DPAMT pages for huge pages is
neither needed nor applicable.

Add a "level" parameter to tdx_pamt_get()/tdx_pamt_put() so the helpers
can skip DPAMT page installation and uninstallation for huge pages.

TD control pages (including S-EPT page table pages) are managed by the TDX
module at the 4KB level. Therefore, pass PG_LEVEL_4K to
tdx_pamt_get()/tdx_pamt_put() for those pages.

This patch is based on earlier work from Kiryl Shutsemau.
Signed-off-by: Yan Zhao <yan.y.zhao@intel.com>
---
 arch/x86/include/asm/tdx.h  |  4 ++--
 arch/x86/kvm/vmx/tdx.c      | 14 +++++++++-----
 arch/x86/virt/vmx/tdx/tdx.c | 12 ++++++------
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/arch/x86/include/asm/tdx.h b/arch/x86/include/asm/tdx.h
index 48efd3eac375..e2790ac2c521 100644
--- a/arch/x86/include/asm/tdx.h
+++ b/arch/x86/include/asm/tdx.h
@@ -135,8 +135,8 @@ static inline void tdx_init_pamt_cache(struct tdx_pamt_cache *cache)
 
 void tdx_free_pamt_cache(struct tdx_pamt_cache *cache);
 int tdx_topup_pamt_cache(struct tdx_pamt_cache *cache, unsigned long npages);
-int tdx_pamt_get(kvm_pfn_t pfn, struct tdx_pamt_cache *cache);
-void tdx_pamt_put(kvm_pfn_t pfn);
+int tdx_pamt_get(kvm_pfn_t pfn, enum pg_level level, struct tdx_pamt_cache *cache);
+void tdx_pamt_put(kvm_pfn_t pfn, enum pg_level level);
 
 int tdx_guest_keyid_alloc(void);
 u32 tdx_get_nr_guest_keyids(void);
diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 1f0dba322640..410c149bf776 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -1710,14 +1710,18 @@ static int tdx_sept_map_nonleaf_spte(struct kvm *kvm, gfn_t gfn,
 	if (!sept_pt)
 		return -EIO;
 
-	ret = tdx_pamt_get(page_to_pfn(sept_pt), &to_tdx(vcpu)->pamt_cache);
+	/*
+	 * An S-EPT page table page is tracked by the TDX module at 4KB level,
+	 * so pass 4KB instead of the target S-EPT level.
+	 */
+	ret = tdx_pamt_get(page_to_pfn(sept_pt), PG_LEVEL_4K, &to_tdx(vcpu)->pamt_cache);
 	if (KVM_BUG_ON(ret, kvm))
 		return ret;
 
 	err = tdh_mem_sept_add(&to_kvm_tdx(kvm)->td, gpa, level, sept_pt,
 			       &entry, &level_state);
 	if (err)
-		tdx_pamt_put(page_to_pfn(sept_pt));
+		tdx_pamt_put(page_to_pfn(sept_pt), PG_LEVEL_4K);
 
 	if (unlikely(tdx_operand_busy(err)))
 		return -EBUSY;
@@ -1745,7 +1749,7 @@ static int tdx_sept_map_leaf_spte(struct kvm *kvm, gfn_t gfn, enum pg_level leve
 
 	WARN_ON_ONCE((new_spte & VMX_EPT_RWX_MASK) != VMX_EPT_RWX_MASK);
 
-	ret = tdx_pamt_get(pfn, &to_tdx(vcpu)->pamt_cache);
+	ret = tdx_pamt_get(pfn, level, &to_tdx(vcpu)->pamt_cache);
 	if (KVM_BUG_ON(ret, kvm))
 		return ret;
 
@@ -1767,7 +1771,7 @@ static int tdx_sept_map_leaf_spte(struct kvm *kvm, gfn_t gfn, enum pg_level leve
 		ret = tdx_mem_page_add(kvm, gfn, level, pfn);
 
 	if (ret)
-		tdx_pamt_put(pfn);
+		tdx_pamt_put(pfn, level);
 
 	return ret;
 }
@@ -1862,7 +1866,7 @@ static int tdx_sept_remove_leaf_spte(struct kvm *kvm, gfn_t gfn,
 		return -EIO;
 
 	tdx_quirk_reset_paddr(PFN_PHYS(pfn), PAGE_SIZE);
-	tdx_pamt_put(pfn);
+	tdx_pamt_put(pfn, level);
 	return 0;
 }
 
diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 33253869e2fd..43f813afc5b1 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -2112,14 +2112,14 @@ static u64 tdh_phymem_pamt_remove(kvm_pfn_t pfn, struct page **pamt_pages)
 static DEFINE_SPINLOCK(dpamt_lock);
 
 /* Bump DPAMT refcount for the given pfn and allocate DPAMT backing if needed. */
-int tdx_pamt_get(kvm_pfn_t pfn, struct tdx_pamt_cache *cache)
+int tdx_pamt_get(kvm_pfn_t pfn, enum pg_level level, struct tdx_pamt_cache *cache)
 {
 	struct page *pamt_pages[TDX_DPAMT_ENTRY_PAGE_CNT];
 	atomic_t *dpamt_refcount;
 	u64 tdx_status;
 	int ret;
 
-	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
+	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo) || level != PG_LEVEL_4K)
 		return 0;
 
 	ret = alloc_pamt_array(pamt_pages, cache);
@@ -2157,13 +2157,13 @@ int tdx_pamt_get(kvm_pfn_t pfn, struct tdx_pamt_cache *cache)
 EXPORT_SYMBOL_FOR_KVM(tdx_pamt_get);
 
 /* Drop DPAMT refcount for the given pfn and free DPAMT backing if needed. */
-void tdx_pamt_put(kvm_pfn_t pfn)
+void tdx_pamt_put(kvm_pfn_t pfn, enum pg_level level)
 {
 	struct page *pamt_pages[TDX_DPAMT_ENTRY_PAGE_CNT] = {};
 	atomic_t *dpamt_refcount;
 	u64 tdx_status;
 
-	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
+	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo) || level != PG_LEVEL_4K)
 		return;
 
 	dpamt_refcount = tdx_find_dpamt_refcount(pfn);
@@ -2243,7 +2243,7 @@ struct page *tdx_alloc_control_page(void)
 	if (!page)
 		return NULL;
 
-	if (tdx_pamt_get(page_to_pfn(page), NULL)) {
+	if (tdx_pamt_get(page_to_pfn(page), PG_LEVEL_4K, NULL)) {
 		__free_page(page);
 		return NULL;
 	}
@@ -2261,7 +2261,7 @@ void tdx_free_control_page(struct page *page)
 	if (!page)
 		return;
 
-	tdx_pamt_put(page_to_pfn(page));
+	tdx_pamt_put(page_to_pfn(page), PG_LEVEL_4K);
 	__free_page(page);
 }
 EXPORT_SYMBOL_FOR_KVM(tdx_free_control_page);
-- 
2.43.2