Re: [PATCH v7 07/24] KVM: arm64: iommu: Shadow host stage-2 page table

[Mostafa Saleh <smostafa@google.com>]

Hi Seb,

On Thu, Jul 23, 2026 at 03:29:37PM +0000, Sebastian Ene wrote:
> On Wed, Jul 15, 2026 at 11:58:48AM +0000, Mostafa Saleh wrote:
> > diff --git a/arch/arm64/kvm/hyp/include/nvhe/iommu.h b/arch/arm64/kvm/hyp/include/nvhe/iommu.h
> > index df3d0cc5d4db..857d7dd2ebc3 100644
> > --- a/arch/arm64/kvm/hyp/include/nvhe/iommu.h
> > +++ b/arch/arm64/kvm/hyp/include/nvhe/iommu.h
> > @@ -3,11 +3,15 @@
> >  #define __ARM64_KVM_NVHE_IOMMU_H__
> >  
> >  #include <asm/kvm_host.h>
> > +#include <asm/kvm_pgtable.h>
> >  
> >  struct pkvm_iommu_ops {
> >  	int (*init)(void);
> > +	int (*host_stage2_idmap)(phys_addr_t start, phys_addr_t end, int prot);
> 
> Hi Mostafa,
> 
> Maybe host_stage2_idmap_locked, since it is invoked while the host
> stage-2 lock is acquired ?

Vincent mentioned in earlier version that it can be called idmap()
also, I believe as this is not a function but a pointer, we do not
have over descirbe the name, but no strong opionion.

> 
> >  };
> >  
> >  int pkvm_iommu_init(void);
> >  
> > +int pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
> > +				 enum kvm_pgtable_prot prot);
> >  #endif /* __ARM64_KVM_NVHE_IOMMU_H__ */
> > diff --git a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
> > index 51b0eb3844a9..99b821b3cf65 100644
> > --- a/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
> > +++ b/arch/arm64/kvm/hyp/include/nvhe/mem_protect.h
> > @@ -59,6 +59,7 @@ int __pkvm_host_test_clear_young_guest(u64 gfn, u64 nr_pages, bool mkold, struct
> >  int __pkvm_host_mkyoung_guest(u64 gfn, struct pkvm_hyp_vcpu *vcpu);
> >  
> >  bool addr_is_memory(phys_addr_t phys);
> > +
> >  int host_stage2_idmap_locked(phys_addr_t addr, u64 size, enum kvm_pgtable_prot prot);
> >  int host_stage2_set_owner_locked(phys_addr_t addr, u64 size, u8 owner_id);
> >  int kvm_host_prepare_stage2(void *pgt_pool_base);
> > diff --git a/arch/arm64/kvm/hyp/nvhe/iommu.c b/arch/arm64/kvm/hyp/nvhe/iommu.c
> > index ef456eff42d2..08009609ec59 100644
> > --- a/arch/arm64/kvm/hyp/nvhe/iommu.c
> > +++ b/arch/arm64/kvm/hyp/nvhe/iommu.c
> > @@ -4,16 +4,137 @@
> >   *
> >   * Copyright (C) 2022 Linaro Ltd.
> >   */
> > +#include <linux/iommu.h>
> > +#include <asm/kvm_pkvm.h>
> > +
> >  #include <nvhe/iommu.h>
> > +#include <nvhe/mem_protect.h>
> > +#include <nvhe/spinlock.h>
> >  
> >  /* Only one set of ops supported */
> >  struct pkvm_iommu_ops *pkvm_iommu_ops;
> >  
> > -int pkvm_iommu_init(void)
> > +/* Protected by host_mmu.lock */
> > +static bool pkvm_idmap_initialized;
> > +
> > +static inline int pkvm_to_iommu_prot(enum kvm_pgtable_prot prot)
> >  {
> > -	/* Keep DMA isolation optional. */
> > -	if (!pkvm_iommu_ops || !pkvm_iommu_ops->init)
> > +	int iommu_prot = 0;
> > +
> > +	if (prot & KVM_PGTABLE_PROT_R)
> > +		iommu_prot |= IOMMU_READ;
> > +	if (prot & KVM_PGTABLE_PROT_W)
> > +		iommu_prot |= IOMMU_WRITE;
> > +
> > +	/* We don't understand that, might be dangerous. */
> > +	WARN_ON(prot & ~PKVM_HOST_MEM_PROT);
> > +	return iommu_prot;
> > +}
> > +
> > +/*
> > + * IOMMU page tables are shadowed and not shared, that is mainly because:
> > + * - Possible inconsistency between IOMMU and CPU features or format.
> > + * - KVM relies on handling in page faults (BBM, lazy mapping).
> > + */
> > +static int __snapshot_host_stage2(const struct kvm_pgtable_visit_ctx *ctx,
> > +				  enum kvm_pgtable_walk_flags visit)
> > +{
> > +	u64 start = ctx->addr;
> > +	u64 block_end = ALIGN_DOWN(ctx->addr, kvm_granule_size(ctx->level)) +
> > +			kvm_granule_size(ctx->level);
> > +	u64 end = min(ctx->end, block_end);
> > +	kvm_pte_t pte = *ctx->ptep;
> > +	bool is_memory = *(bool *)ctx->arg;
> > +	int prot;
> > +
> > +	/*
> > +	 * Keep annotated PTEs unmapped, and map everything else even lazily
> > +	 * mapped PTEs(0), as the IOMMU can't handle page faults.
> > +	 * That maps the whole address space which can be large, but that doesn't
> > +	 * use a lot of memory as it will be mostly large block (1 GB with 4kb pages)
> > +	 */
> > +	if (pte && !kvm_pte_valid(pte))
> >  		return 0;
> 
> >  
> > -	return pkvm_iommu_ops->init();
> > +	if (kvm_pte_valid(pte))
> > +		prot = pkvm_to_iommu_prot(kvm_pgtable_stage2_pte_prot(pte));
> > +	else
> > +		prot = IOMMU_READ | IOMMU_WRITE;
> 
> If it's invalid in the host stage-2 why do you make it read/write in the IOMMU
> ? Shouldn't it be an invalid descriptor in the IOMMU pagetables as well
> ?
> 

These are invalid and empty PTEs as a few lines up we return early for
invalid non-empty ones. For the empty PTEs, they are owned by the host
but lazily mapped. The IOMMU page table need to eagrly map those as it
can not handle page faults like the CPU.

> > +
> > +	if (!is_memory)
> > +		prot |= IOMMU_MMIO;
> > +
> > +	return pkvm_iommu_ops->host_stage2_idmap(start, end, prot);
> > +}
> > +
> > +static int pkvm_iommu_snapshot_host_stage2(void)
> > +{
> > +	struct kvm_pgtable *pgt = &host_mmu.pgt;
> > +	bool is_memory;
> > +	struct kvm_pgtable_walker walker = {
> > +		.cb	= __snapshot_host_stage2,
> > +		.flags	= KVM_PGTABLE_WALK_LEAF,
> > +		.arg	= &is_memory,
> > +	};
> > +	int ret = 0, i;
> > +	u64 start = 0;
> > +
> > +	hyp_spin_lock(&host_mmu.lock);
> > +	for (i = 0; i < hyp_memblock_nr; i++) {
> > +		struct memblock_region *reg = &hyp_memory[i];
> > +
> > +		if (start < reg->base) {
> > +			is_memory = false;
> > +			ret = kvm_pgtable_walk(pgt, start, reg->base - start, &walker);
> > +			if (ret)
> > +				goto out_unlock;
> > +		}
> > +
> > +		is_memory = true;
> > +		ret = kvm_pgtable_walk(pgt, reg->base, reg->size, &walker);
> > +		if (ret)
> > +			goto out_unlock;
> > +
> > +		start = reg->base + reg->size;
> > +	}
> > +
> > +	if (start < BIT(pgt->ia_bits)) {
> > +		is_memory = false;
> > +		ret = kvm_pgtable_walk(pgt, start, BIT(pgt->ia_bits) - start, &walker);
> > +		if (ret)
> > +			goto out_unlock;
> > +	}
> > +
> > +	pkvm_idmap_initialized = true;
> > +
> > +out_unlock:
> > +	hyp_spin_unlock(&host_mmu.lock);
> > +	return ret;
> > +}
> > +
> > +int pkvm_iommu_init(void)
> > +{
> > +	int ret;
> > +
> > +	/* Keep DMA isolation optional. */
> > +	if (!pkvm_iommu_ops || !pkvm_iommu_ops->init ||
> > +	    !pkvm_iommu_ops->host_stage2_idmap)
> > +		return 0;
> 
> Shouldn't this scream or return not supported ? 
> 

This was the case in the earlier versions, but I figured it would
regress so many setups, so I made DMA isolation optinal for now, I
think long term this can be a cmdline option. But t's up to the
maintainers how do they want to support this. I am happy to bring
the error return back.

> > +
> > +	ret = pkvm_iommu_ops->init();
> > +	if (ret)
> > +		return ret;
> > +
> > +	return pkvm_iommu_snapshot_host_stage2();
> 
> If pkvm_iommu_snapshot_host_stage2() fails would you have to provide a
> ->deinit() corresponding callback ?
> 

Good point, I tried to avoid having a deinit by making the IOMMU the
last thing done during setup as I did not want to make the patches
bigger.

Maybe WARN_ON() is enough for now? I need to think more about it.

Thanks,
Mostafa