[RFC PATCH v2 22/23] KVM: TDX: Handle Dynamic PAMT on page split

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

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

Page demote from 2M to 4k requires an additional PAMT page pair to cover
the 2M range that now mapped with 4k.

EPT page also has to be covered in PAMT_4K.

Allocate both from pre-allocated split PAMT pool.

Signed-off-by: Kirill A. Shutemov <kirill.shutemov@linux.intel.com>
Signed-off-by: Yan Zhao <yan.y.zhao@intel.com>
---
RFC v2:
- Pulled from
  git://git.kernel.org/pub/scm/linux/kernel/git/kas/linux.git tdx/dpamt-huge.
- Rebased on top of TDX huge page RFC v2 (Yan).
---
 arch/x86/include/asm/tdx.h  |  4 ++++
 arch/x86/kvm/vmx/tdx.c      | 28 ++++++++++++++++++++++++----
 arch/x86/virt/vmx/tdx/tdx.c | 11 +++++++----
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/arch/x86/include/asm/tdx.h b/arch/x86/include/asm/tdx.h
index 2e529f0c578a..da317981e95a 100644
--- a/arch/x86/include/asm/tdx.h
+++ b/arch/x86/include/asm/tdx.h
@@ -123,6 +123,10 @@ u32 tdx_get_nr_guest_keyids(void);
 void tdx_guest_keyid_free(unsigned int keyid);
 
 int tdx_nr_pamt_pages(void);
+atomic_t *tdx_get_pamt_refcount(unsigned long hpa);
+int tdx_alloc_pamt_pages(struct list_head *pamt_pages,
+			 struct page *(alloc)(void *data), void *data);
+void tdx_free_pamt_pages(struct list_head *pamt_pages);
 int tdx_pamt_get(struct page *page, enum pg_level level,
 		 struct page *(alloc)(void *data), void *data);
 void tdx_pamt_put(struct page *page, enum pg_level level);
diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c
index 9d24a1a86a23..6e061d659639 100644
--- a/arch/x86/kvm/vmx/tdx.c
+++ b/arch/x86/kvm/vmx/tdx.c
@@ -1915,28 +1915,48 @@ static int tdx_sept_free_private_spt(struct kvm *kvm, gfn_t gfn,
 	return 0;
 }
 
+static struct page *tdx_alloc_pamt_page_split(void *data)
+{
+	struct kvm *kvm = data;
+	void *p;
+
+	p = kvm_mmu_memory_cache_alloc(&kvm->arch.pamt_page_cache);
+	return virt_to_page(p);
+}
+
 static int tdx_spte_demote_private_spte(struct kvm *kvm, gfn_t gfn,
-					enum pg_level level, struct page *page)
+					enum pg_level level, struct page *page,
+					kvm_pfn_t pfn_for_gfn)
 {
 	int tdx_level = pg_level_to_tdx_sept_level(level);
+	hpa_t hpa = pfn_to_hpa(pfn_for_gfn);
 	struct kvm_tdx *kvm_tdx = to_kvm_tdx(kvm);
 	gpa_t gpa = gfn_to_gpa(gfn);
 	u64 err, entry, level_state;
+	LIST_HEAD(pamt_pages);
+
+	tdx_pamt_get(page, PG_LEVEL_4K, tdx_alloc_pamt_page_split, kvm);
+	tdx_alloc_pamt_pages(&pamt_pages, tdx_alloc_pamt_page_split, kvm);
 
 	err = tdh_mem_page_demote(&kvm_tdx->td, gpa, tdx_level, page,
-				  NULL, &entry, &level_state);
+				  &pamt_pages, &entry, &level_state);
 
 	if (unlikely(tdx_operand_busy(err))) {
 		tdx_no_vcpus_enter_start(kvm);
 		err = tdh_mem_page_demote(&kvm_tdx->td, gpa, tdx_level, page,
-					  NULL, &entry, &level_state);
+					  &pamt_pages, &entry, &level_state);
 		tdx_no_vcpus_enter_stop(kvm);
 	}
 
 	if (KVM_BUG_ON(err, kvm)) {
+		tdx_free_pamt_pages(&pamt_pages);
+		tdx_pamt_put(page, PG_LEVEL_4K);
 		pr_tdx_error_2(TDH_MEM_PAGE_DEMOTE, err, entry, level_state);
 		return -EIO;
 	}
+
+	if (tdx_supports_dynamic_pamt(tdx_sysinfo))
+		atomic_set(tdx_get_pamt_refcount(hpa), PTRS_PER_PMD);
 	return 0;
 }
 
@@ -1963,7 +1983,7 @@ static int tdx_sept_split_private_spt(struct kvm *kvm, gfn_t gfn, enum pg_level
 
 	tdx_track(kvm);
 
-	return tdx_spte_demote_private_spte(kvm, gfn, level, page);
+	return tdx_spte_demote_private_spte(kvm, gfn, level, page, pfn_for_gfn);
 }
 
 static int tdx_sept_remove_private_spte(struct kvm *kvm, gfn_t gfn,
diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 50f9d49f1c91..dbbddd00ec60 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -188,10 +188,11 @@ int tdx_cpu_enable(void)
 }
 EXPORT_SYMBOL_GPL(tdx_cpu_enable);
 
-static atomic_t *tdx_get_pamt_refcount(unsigned long hpa)
+atomic_t *tdx_get_pamt_refcount(unsigned long hpa)
 {
 	return &pamt_refcounts[hpa / PMD_SIZE];
 }
+EXPORT_SYMBOL_GPL(tdx_get_pamt_refcount);
 
 static int pamt_refcount_populate(pte_t *pte, unsigned long addr, void *data)
 {
@@ -2151,7 +2152,7 @@ static u64 tdh_phymem_pamt_remove(unsigned long hpa,
 
 static DEFINE_SPINLOCK(pamt_lock);
 
-static void tdx_free_pamt_pages(struct list_head *pamt_pages)
+void tdx_free_pamt_pages(struct list_head *pamt_pages)
 {
 	struct page *page;
 
@@ -2160,9 +2161,10 @@ static void tdx_free_pamt_pages(struct list_head *pamt_pages)
 		__free_page(page);
 	}
 }
+EXPORT_SYMBOL_GPL(tdx_free_pamt_pages);
 
-static int tdx_alloc_pamt_pages(struct list_head *pamt_pages,
-				 struct page *(alloc)(void *data), void *data)
+int tdx_alloc_pamt_pages(struct list_head *pamt_pages,
+			 struct page *(alloc)(void *data), void *data)
 {
 	for (int i = 0; i < tdx_nr_pamt_pages(); i++) {
 		struct page *page;
@@ -2180,6 +2182,7 @@ static int tdx_alloc_pamt_pages(struct list_head *pamt_pages,
 	tdx_free_pamt_pages(pamt_pages);
 	return -ENOMEM;
 }
+EXPORT_SYMBOL_GPL(tdx_alloc_pamt_pages);
 
 static int tdx_pamt_add(atomic_t *pamt_refcount, unsigned long hpa,
 			struct list_head *pamt_pages)
-- 
2.43.2