ANDROID: KVM: arm64: Add flags for hyp memcache topup and free

So far the only flag is to enable or not the kmemcg accounting, this
will enable to use those function to refill the hyp heap allocator.

Bug: 357781595
Bug: 273748186
Change-Id: I35d701d23f5d66ab06aed50c62b3b5c6cbf46e14
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -138,8 +138,12 @@ static inline void __free_hyp_memcache(struct kvm_hyp_memcache *mc,
 		free_fn(pop_hyp_memcache(mc, to_va), arg);
 }
 
-void free_hyp_memcache(struct kvm_hyp_memcache *mc);
-int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages);
+#define HYP_MEMCACHE_ACCOUNT_KMEMCG BIT(1)
+
+void free_hyp_memcache(struct kvm_hyp_memcache *mc,
+		       unsigned long flags);
+int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages,
+		       unsigned long flags);
 
 struct kvm_vmid {
 	atomic64_t id;

arch/arm64/kvm/arm.c

@@ -528,7 +528,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);
+		free_hyp_memcache(&vcpu->arch.pkvm_memcache, 0);
 	else
 		kvm_mmu_free_memory_cache(&vcpu->arch.mmu_page_cache);
 

arch/arm64/kvm/hypercalls.c

@@ -380,7 +380,8 @@ int kvm_smccc_call_handler(struct kvm_vcpu *vcpu)
 	case ARM_SMCCC_VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID:
 		if (kvm_vm_is_protected(vcpu->kvm) &&
 		    !topup_hyp_memcache(&vcpu->arch.pkvm_memcache,
-					kvm_mmu_cache_min_pages(&vcpu->kvm->arch.mmu)))
+					kvm_mmu_cache_min_pages(&vcpu->kvm->arch.mmu),
+					HYP_MEMCACHE_ACCOUNT_KMEMCG))
 			val[0] = SMCCC_RET_SUCCESS;
 		break;
 	case ARM_SMCCC_VENDOR_HYP_KVM_MEM_RELINQUISH_FUNC_ID:

arch/arm64/kvm/mmu.c

@@ -1176,30 +1176,37 @@ 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 *flags)
 {
 	free_page((unsigned long)addr);
 }
 
-static void *hyp_mc_alloc_fn(void *unused)
+static void *hyp_mc_alloc_fn(void *flags)
 {
-	return (void *)__get_free_page(GFP_KERNEL_ACCOUNT);
+	unsigned long __flags = (unsigned long)flags;
+	gfp_t gfp_mask;
+
+	gfp_mask = __flags & HYP_MEMCACHE_ACCOUNT_KMEMCG ?
+		   GFP_KERNEL_ACCOUNT : GFP_KERNEL;
+
+	return (void *)__get_free_page(gfp_mask);
 }
 
-void free_hyp_memcache(struct kvm_hyp_memcache *mc)
+void free_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long flags)
 {
 	if (is_protected_kvm_enabled())
 		__free_hyp_memcache(mc, hyp_mc_free_fn,
-				    kvm_host_va, NULL);
+				    kvm_host_va, (void *)flags);
 }
 
-int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages)
+int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages,
+		       unsigned long flags)
 {
 	if (!is_protected_kvm_enabled())
 		return 0;
 
 	return __topup_hyp_memcache(mc, min_pages, hyp_mc_alloc_fn,
-				    kvm_host_pa, NULL);
+				    kvm_host_pa, (void *)flags);
 }
 
 /**
@@ -1642,7 +1649,8 @@ static int pkvm_mem_abort(struct kvm_vcpu *vcpu, phys_addr_t fault_ipa,
 	u64 pfn;
 	int ret;
 
-	ret = topup_hyp_memcache(hyp_memcache, kvm_mmu_cache_min_pages(mmu));
+	ret = topup_hyp_memcache(hyp_memcache, kvm_mmu_cache_min_pages(mmu),
+				 HYP_MEMCACHE_ACCOUNT_KMEMCG);
 	if (ret)
 		return -ENOMEM;
 

arch/arm64/kvm/pkvm.c

@@ -264,7 +264,7 @@ static void __pkvm_destroy_hyp_vm(struct kvm *host_kvm)
 
 out_free:
 	host_kvm->arch.pkvm.handle = 0;
-	free_hyp_memcache(&host_kvm->arch.pkvm.teardown_mc);
+	free_hyp_memcache(&host_kvm->arch.pkvm.teardown_mc, 0);
 }
 
 static void __pkvm_vcpu_hyp_created(struct kvm_vcpu *vcpu)