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.

Bug: 278749606
Bug: 222044477
Change-Id: I2a45ce65c5ce9cf96aabd1b66d6f83ffe4808a0c
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index c711db7..3a7298a 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -139,6 +139,7 @@ static inline void __free_hyp_memcache(struct kvm_hyp_memcache *mc,
 }
 
 #define HYP_MEMCACHE_ACCOUNT_KMEMCG BIT(1)
+#define HYP_MEMCACHE_ACCOUNT_STAGE2 BIT(2)
 
 void free_hyp_memcache(struct kvm_hyp_memcache *mc);
 int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages);
@@ -151,7 +152,7 @@ static inline void init_hyp_memcache(struct kvm_hyp_memcache *mc)
 static inline void init_hyp_stage2_memcache(struct kvm_hyp_memcache *mc)
 {
 	init_hyp_memcache(mc);
-	mc->flags = HYP_MEMCACHE_ACCOUNT_KMEMCG;
+	mc->flags = HYP_MEMCACHE_ACCOUNT_KMEMCG | HYP_MEMCACHE_ACCOUNT_STAGE2;
 }
 
 struct kvm_vmid {
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 99934e9..81799fd 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -1178,6 +1178,12 @@ void kvm_free_stage2_pgd(struct kvm_s2_mmu *mmu)
 
 static void hyp_mc_free_fn(void *addr, void *flags)
 {
+	if (!addr)
+		return;
+
+	if ((unsigned long)flags & HYP_MEMCACHE_ACCOUNT_STAGE2)
+		kvm_account_pgtable_pages(addr, -1);
+
 	free_page((unsigned long)addr);
 }
 
@@ -1185,11 +1191,17 @@ static void *hyp_mc_alloc_fn(void *flags)
 {
 	unsigned long __flags = (unsigned long)flags;
 	gfp_t gfp_mask;
+	void *addr;
 
 	gfp_mask = __flags & HYP_MEMCACHE_ACCOUNT_KMEMCG ?
 		   GFP_KERNEL_ACCOUNT : GFP_KERNEL;
 
-	return (void *)__get_free_page(gfp_mask);
+	addr = (void *)__get_free_page(gfp_mask);
+
+	if (addr && __flags & HYP_MEMCACHE_ACCOUNT_STAGE2)
+		kvm_account_pgtable_pages(addr, 1);
+
+	return addr;
 }
 
 void free_hyp_memcache(struct kvm_hyp_memcache *mc)
diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index d207a69..a4f7e7f 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -313,6 +313,8 @@ static int __pkvm_create_hyp_vm(struct kvm *host_kvm)
 		__pkvm_vcpu_hyp_created(host_vcpu);
 	}
 
+	kvm_account_pgtable_pages(pgd, pgd_sz >> PAGE_SHIFT);
+
 	return 0;
 
 destroy_vm: