ANDROID: KVM: arm64: iommu: Add hyp_pool for guest pvIOMMU

Hyp pool would be used to cache guest allocations for the IOMMU, so
when they are freed they can be used again.

Also add a memcache for teardown of the pool.

Bug: 357781595
Bug: 348382247
Bug: 236685427
Change-Id: I43b0e097017efef4147ee5e94c7a77b4a251acb8
Signed-off-by: Mostafa Saleh <smostafa@google.com>

arch/arm64/include/asm/kvm_host.h

@@ -295,6 +295,7 @@ struct kvm_protected_vm {
 	pkvm_handle_t handle;
 	struct kvm_hyp_memcache stage2_teardown_mc;
 	struct rb_root_cached pinned_pages;
+	struct kvm_hyp_memcache teardown_iommu_mc;
 	gpa_t pvmfw_load_addr;
 	bool enabled;
 };

arch/arm64/kvm/hyp/include/nvhe/mem_protect.h

@@ -103,7 +103,7 @@ int reclaim_hyp_pool(struct hyp_pool *pool, struct kvm_hyp_memcache *host_mc,
 		     int nr_pages);
 
 void destroy_hyp_vm_pgt(struct pkvm_hyp_vm *vm);
-void drain_hyp_pool(struct pkvm_hyp_vm *vm, struct kvm_hyp_memcache *mc);
+void drain_hyp_pool(struct hyp_pool *pool, struct kvm_hyp_memcache *mc);
 
 int module_change_host_page_prot(u64 pfn, enum kvm_pgtable_prot prot, u64 nr_pages, bool update_iommu);
 

arch/arm64/kvm/hyp/include/nvhe/pkvm.h

@@ -67,6 +67,7 @@ struct pkvm_hyp_vm {
 
 	/* pvIOMMUs attached. */
 	struct list_head pviommus;
+	struct hyp_pool iommu_pool;
 
 	/* Primary vCPU pending entry to the pvmfw */
 	struct pkvm_hyp_vcpu *pvmfw_entry_vcpu;

arch/arm64/kvm/hyp/nvhe/iommu/pviommu-host.c

@@ -98,6 +98,11 @@ int pkvm_pviommu_add_vsid(struct kvm *host_kvm, int pviommu,
 int pkvm_pviommu_finalise(struct pkvm_hyp_vm *hyp_vm)
 {
 	int i;
+	int ret;
+
+	ret = hyp_pool_init_empty(&hyp_vm->iommu_pool, 64);
+	if (ret)
+		return ret;
 
 	hyp_spin_lock(&host_pviommu_lock);
 	INIT_LIST_HEAD(&hyp_vm->pviommus);

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

@@ -2362,9 +2362,9 @@ void destroy_hyp_vm_pgt(struct pkvm_hyp_vm *vm)
 	guest_unlock_component(vm);
 }
 
-void drain_hyp_pool(struct pkvm_hyp_vm *vm, struct kvm_hyp_memcache *mc)
+void drain_hyp_pool(struct hyp_pool *pool, struct kvm_hyp_memcache *mc)
 {
-	WARN_ON(reclaim_hyp_pool(&vm->pool, mc, INT_MAX) != -ENOMEM);
+	WARN_ON(reclaim_hyp_pool(pool, mc, INT_MAX) != -ENOMEM);
 }
 
 int __pkvm_host_reclaim_page(struct pkvm_hyp_vm *vm, u64 pfn, u64 ipa, u8 order)

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

@@ -373,7 +373,7 @@ int __pkvm_reclaim_dying_guest_page(pkvm_handle_t handle, u64 pfn, u64 gfn, u8 o
 	if (ret)
 		goto unlock;
 
-	drain_hyp_pool(hyp_vm, &hyp_vm->host_kvm->arch.pkvm.stage2_teardown_mc);
+	drain_hyp_pool(&hyp_vm->pool, &hyp_vm->host_kvm->arch.pkvm.stage2_teardown_mc);
 unlock:
 	hyp_read_unlock(&vm_table_lock);
 
@@ -1011,6 +1011,13 @@ int __pkvm_finalize_teardown_vm(pkvm_handle_t handle)
 
 	pkvm_pviommu_teardown(hyp_vm);
 
+	/*
+	 * At this point all page tables are destroyed and should be pushed to the pool
+	 * the only place that might still have memory is the mc, which would be drained
+	 * from host as it hasn't been donated yet.
+	 */
+	drain_hyp_pool(&hyp_vm->iommu_pool, &host_kvm->arch.pkvm.teardown_iommu_mc);
+
 	/*
 	 * At this point, the VM has been detached from the VM table and
 	 * has a refcount of 0 so we're free to tear it down without
@@ -1019,7 +1026,7 @@ int __pkvm_finalize_teardown_vm(pkvm_handle_t handle)
 
 	mc = &host_kvm->arch.pkvm.stage2_teardown_mc;
 	destroy_hyp_vm_pgt(hyp_vm);
-	drain_hyp_pool(hyp_vm, mc);
+	drain_hyp_pool(&hyp_vm->pool, mc);
 	unpin_host_vcpus(hyp_vm->vcpus, hyp_vm->kvm.created_vcpus);
 
 	/* Push the metadata pages to the teardown memcache */

arch/arm64/kvm/pkvm.c

@@ -425,6 +425,8 @@ out_free:
 		     &host_kvm->stat.protected_hyp_mem);
 	free_hyp_memcache(&host_kvm->arch.pkvm.stage2_teardown_mc);
 
+	kvm_iommu_guest_free_mc(&host_kvm->arch.pkvm.teardown_iommu_mc);
+
 	kvm_for_each_vcpu(idx, host_vcpu, host_kvm) {
 		struct kvm_hyp_req *hyp_reqs = host_vcpu->arch.hyp_reqs;