From: "Kirill A. Shutemov" <kirill.shutemov@linux.intel.com>

Invoke tdx_pamt_{get/put}() to add/remove Dynamic PAMT page pair for guest
private memory only when the S-EPT mapping size is 4KB.

When the mapping size is greater than 4KB, static PAMT pages are used. No
need to install/uninstall extra PAMT pages dynamically.

Signed-off-by: Kirill A. Shutemov <kirill.shutemov@linux.intel.com>
[Yan: Move level checking to callers of tdx_pamt_{get/put}()]
Signed-off-by: Yan Zhao <yan.y.zhao@intel.com>
---
v3:
- new patch

  Checking for 4KB level was previously done inside tdx_pamt_{get/put}() in
  DPAMT v2 [1].

  Move the checking to callers of tdx_pamt_{get/put}() in KVM to avoid
  introducing an extra "level" parameter to tdx_pamt_{get/put}(). This is
  also because the callers that could have level > 4KB are limited in KVM,
  i.e., only inside tdx_sept_{set/remove}_private_spte().

[1] https://lore.kernel.org/all/20250609191340.2051741-5-kirill.shutemov@linux.intel.com
---
 arch/x86/kvm/vmx/tdx.c | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 712aaa3d45b7..c1dc1aaae49d 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -1722,9 +1722,11 @@ static int tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
 	WARN_ON_ONCE(!is_shadow_present_pte(mirror_spte) ||
 		     (mirror_spte & VMX_EPT_RWX_MASK) != VMX_EPT_RWX_MASK);
 
-	ret = tdx_pamt_get(page, &tdx->prealloc);
-	if (ret)
-		return ret;
+	if (level == PG_LEVEL_4K) {
+		ret = tdx_pamt_get(page, &tdx->prealloc);
+		if (ret)
+			return ret;
+	}
 
 	/*
 	 * Ensure pre_fault_allowed is read by kvm_arch_vcpu_pre_fault_memory()
@@ -1743,7 +1745,7 @@ static int tdx_sept_set_private_spte(struct kvm *kvm, gfn_t gfn,
 	else
 		ret = tdx_mem_page_add(kvm, gfn, level, pfn);
 
-	if (ret)
+	if (ret && level == PG_LEVEL_4K)
 		tdx_pamt_put(page);
 
 	return ret;
@@ -1911,7 +1913,9 @@ static void tdx_sept_remove_private_spte(struct kvm *kvm, gfn_t gfn,
 
 	tdx_quirk_reset_folio(folio, folio_page_idx(folio, page),
 			      KVM_PAGES_PER_HPAGE(level));
-	tdx_pamt_put(page);
+
+	if (level == PG_LEVEL_4K)
+		tdx_pamt_put(page);
 }
 
 /*
-- 
2.43.2