ANDROID: KVM: arm64: count KVM s2 mmu usage in nVHE protected mode

When using the nVHE protected mode, the stage-2 page tables are handled by
the hypervisor, but are backed by memory donated by the host. That memory
is accounted during the donation (add to the vCPUs hyp_memcache) under
secondary pagetable stats.

On VM teardown, those pages are mixed with others in the teardown_mc, so use
a separated teardown_stage2_mc to deduct them from accounting after
reclaim.

Bug: 222044477
Change-Id: I2a45ce65c5ce9cf96aabd1b66d6f83ffe4808a0c
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
Signed-off-by: Quentin Perret <qperret@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -128,6 +128,7 @@ static inline void __free_hyp_memcache(struct kvm_hyp_memcache *mc,
 }
 
 void free_hyp_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm);
+void free_hyp_stage2_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm);
 int topup_hyp_memcache(struct kvm_vcpu *vcpu);
 
 struct kvm_vmid {
@@ -183,6 +184,7 @@ typedef unsigned int pkvm_handle_t;
 struct kvm_protected_vm {
 	pkvm_handle_t handle;
 	struct kvm_hyp_memcache teardown_mc;
+	struct kvm_hyp_memcache teardown_stage2_mc;
 	struct rb_root pinned_pages;
 	gpa_t pvmfw_load_addr;
 	bool enabled;

arch/arm64/kvm/arm.c

@@ -463,7 +463,7 @@ void kvm_arch_vcpu_destroy(struct kvm_vcpu *vcpu)
 		static_branch_dec(&userspace_irqchip_in_use);
 
 	if (is_protected_kvm_enabled())
-		free_hyp_memcache(&vcpu->arch.pkvm_memcache, vcpu->kvm);
+		free_hyp_stage2_memcache(&vcpu->arch.pkvm_memcache, vcpu->kvm);
 	else
 		kvm_mmu_free_memory_cache(&vcpu->arch.mmu_page_cache);
 

arch/arm64/kvm/hyp/nvhe/pkvm.c

@@ -798,9 +798,9 @@ teardown_donated_memory(struct kvm_hyp_memcache *mc, void *addr, size_t size)
 
 int __pkvm_teardown_vm(pkvm_handle_t handle)
 {
+	struct kvm_hyp_memcache *mc, *stage2_mc;
 	size_t vm_size, last_ran_size;
 	int __percpu *last_vcpu_ran;
-	struct kvm_hyp_memcache *mc;
 	struct pkvm_hyp_vm *hyp_vm;
 	unsigned int idx;
 	int err;
@@ -822,9 +822,11 @@ int __pkvm_teardown_vm(pkvm_handle_t handle)
 	remove_vm_table_entry(handle);
 	hyp_spin_unlock(&vm_table_lock);
 
-	/* Reclaim guest pages (including page-table pages) */
 	mc = &hyp_vm->host_kvm->arch.pkvm.teardown_mc;
-	reclaim_guest_pages(hyp_vm, mc);
+	stage2_mc = &hyp_vm->host_kvm->arch.pkvm.teardown_stage2_mc;
+
+	/* Reclaim guest pages (including page-table pages) */
+	reclaim_guest_pages(hyp_vm, stage2_mc);
 	unpin_host_vcpus(hyp_vm->vcpus, hyp_vm->nr_vcpus);
 
 	/* Push the metadata pages to the teardown memcache */
@@ -838,7 +840,7 @@ int __pkvm_teardown_vm(pkvm_handle_t handle)
 		vcpu_mc = &hyp_vcpu->vcpu.arch.pkvm_memcache;
 		while (vcpu_mc->nr_pages) {
 			addr = pop_hyp_memcache(vcpu_mc, hyp_phys_to_virt);
-			push_hyp_memcache(mc, addr, hyp_virt_to_phys);
+			push_hyp_memcache(stage2_mc, addr, hyp_virt_to_phys);
 			unmap_donated_memory_noclear(addr, PAGE_SIZE);
 		}
 

arch/arm64/kvm/mmu.c

@@ -861,14 +861,24 @@ void kvm_free_stage2_pgd(struct kvm_s2_mmu *mmu)
 	}
 }
 
-static void hyp_mc_free_fn(void *addr, void *unused)
+static void hyp_mc_free_fn(void *addr, void *args)
 {
+	bool account_stage2 = (bool)args;
+
+	if (account_stage2)
+		kvm_account_pgtable_pages(addr, -1);
+
 	free_page((unsigned long)addr);
 }
 
 static void *hyp_mc_alloc_fn(void *unused)
 {
-	return (void *)__get_free_page(GFP_KERNEL_ACCOUNT);
+	void *addr = (void *)__get_free_page(GFP_KERNEL_ACCOUNT);
+
+	if (addr)
+		kvm_account_pgtable_pages(addr, 1);
+
+	return addr;
 }
 
 static void account_hyp_memcache(struct kvm_hyp_memcache *mc,
@@ -889,7 +899,9 @@ static void account_hyp_memcache(struct kvm_hyp_memcache *mc,
 	}
 }
 
-void free_hyp_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm)
+static void __free_account_hyp_memcache(struct kvm_hyp_memcache *mc,
+					struct kvm *kvm,
+					bool account_stage2)
 {
 	unsigned long prev_nr_pages;
 
@@ -897,10 +909,27 @@ void free_hyp_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm)
 		return;
 
 	prev_nr_pages = mc->nr_pages;
-	__free_hyp_memcache(mc, hyp_mc_free_fn, kvm_host_va, NULL);
+	__free_hyp_memcache(mc, hyp_mc_free_fn, kvm_host_va,
+			    (void *)account_stage2);
 	account_hyp_memcache(mc, prev_nr_pages, kvm);
 }
 
+void free_hyp_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm)
+{
+	__free_account_hyp_memcache(mc, kvm, false);
+}
+
+/*
+ * All pages donated to the hypervisor through kvm_hyp_memcache are for the
+ * stage-2 page table. However, kvm_hyp_memcache is also a vehicule to retrieve
+ * meta-data from the hypervisor, hence the need for a stage2 specific free
+ * function.
+ */
+void free_hyp_stage2_memcache(struct kvm_hyp_memcache *mc, struct kvm *kvm)
+{
+	__free_account_hyp_memcache(mc, kvm, true);
+}
+
 int topup_hyp_memcache(struct kvm_vcpu *vcpu)
 {
 	struct kvm_hyp_memcache *mc = &vcpu->arch.pkvm_memcache;

arch/arm64/kvm/pkvm.c

@@ -195,6 +195,7 @@ static int __pkvm_create_hyp_vm(struct kvm *host_kvm)
 	}
 
 	atomic64_set(&host_kvm->stat.protected_hyp_mem, total_sz);
+	kvm_account_pgtable_pages(pgd, pgd_sz >> PAGE_SHIFT);
 
 	return 0;
 
@@ -235,6 +236,8 @@ void pkvm_destroy_hyp_vm(struct kvm *host_kvm)
 
 	host_kvm->arch.pkvm.handle = 0;
 	free_hyp_memcache(&host_kvm->arch.pkvm.teardown_mc, host_kvm);
+	free_hyp_stage2_memcache(&host_kvm->arch.pkvm.teardown_stage2_mc,
+				 host_kvm);
 
 	node = rb_first(&host_kvm->arch.pkvm.pinned_pages);
 	while (node) {