[RFC PATCH v3 0/5] iommu/arm-smmu-v3: Use pinned KVM VMID for stage 2

[RFC PATCH v3 0/5] iommu/arm-smmu-v3: Use pinned KVM VMID for stage 2

[Shameer Kolothum]

Hi All,

Changes from RFCv2[0]

-Based on feedback received, removed generic KVM APIs for pinned_vmid_
 {get,put}() and are now implemented directly by KVM ARM.
-struct kvm * is now associated with struct iommufd_device and the association
 is established when idev is first allocated during iommufd_device_bind().
-This is now based on Nicolin's recent series[1], which decouples the vmid
 from S2 parent domains and is assigned during viommu/vSMMU alloc time.
-Added additional checks to make sure when BTM is enabled, kvm is mandatory.
-At the moment, kvm_get_kvm_safe()/kvm_put_kvm() is used inside the smmuv3
 driver and not in iommufd as iommufd is not a user of kvm and is just
 passed through.

ToDo:
-May be we need to check both KVM and SMMUv3 has the same number of VMID
bits supported. Not sure we have systems like that out there.

Why we require this(Thanks to Jean for the nice write-up below):

---
* When enabling BTM in the SMMU, all TLB invalidations to the
  inner-shareable domain issued by the CPU are taken into account by the
  SMMU. That includes for example the TLBI IPAS2E1IS from
  __kvm_tlb_flush_vmid_range().

* BTM is enabled globally in the SMMU CR2 register. If we enable BTM for
  host SVA, then it also affects KVM.

* Stage-1 TLB entries in the SMMU have a bit (ASET) saying "this entry
  is private and does not participate in BTM", which we set for private
  SMMU address spaces.

  Annoyingly, the stage-2 TLB entries do not have it. With BTM all VMIDs
  are shared between CPU and SMMU.

* So, if the SMMU driver allocates VMID privately and we enable BTM, then
  CPU invalidations will remove unrelated SMMU TLB entries. Instead, the
  SMMU driver needs to coordinate with KVM on VMID allocation.

* Private stage-2 address spaces in the SMMU would need to allocate VMIDs
  that aren't used by KVM, but that's not a use-case at the moment:

  - For assigning devices to a host process or to a VM, we use private
    stage-1 mappings. stage-2 will be used to enable nesting translation,
    and will typically mirror the KVM stage-2 since it pins the guest
    address space.

  - If the SMMU doesn't support stage-1, the driver falls back to stage-2
    for private address spaces. For such an implementation we disable BTM.
---

This is sanity tested on a HiSilicon platform and the complete branch is
available here[2].

Please take a look and let me know your feedback.

Thanks,
Shameer
[0] https://lore.kernel.org/kvmarm/20240208151837.35068-1-shameerali.kolothum.thodi@huawei.com/
[1] https://lore.kernel.org/linux-iommu/cover.1741150594.git.nicolinc@nvidia.com/
[2] https://github.com/hisilicon/kernel-dev/tree/smmuv3_vmid-v1-with-rmr-vmid-v3-ext

Jean-Philippe Brucker (1):
  iommu/arm-smmu-v3: Enable broadcast TLB maintenance

Shameer Kolothum (4):
  KVM: arm64: Introduce support to pin VMIDs
  iommufd/device: Associate a kvm pointer to iommufd_device
  iommu/arm-smmu-v3-iommufd: Pass in kvm pointer to viommu_alloc
  iommu/arm-smmu-v3-iommufd: Use KVM VMID for s2 stage

 arch/arm64/include/asm/kvm_host.h             |  3 +
 arch/arm64/kvm/vmid.c                         | 76 ++++++++++++++++++-
 .../arm/arm-smmu-v3/arm-smmu-v3-iommufd.c     | 54 ++++++++++++-
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c   | 24 +++++-
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h   |  3 +
 drivers/iommu/iommufd/device.c                |  5 +-
 drivers/iommu/iommufd/iommufd_private.h       |  2 +
 drivers/iommu/iommufd/viommu.c                |  3 +-
 drivers/vfio/iommufd.c                        |  2 +-
 include/linux/iommu.h                         |  4 +-
 include/linux/iommufd.h                       |  4 +-
 11 files changed, 168 insertions(+), 12 deletions(-)

-- 
2.47.0

[RFC PATCH v3 1/5] KVM: arm64: Introduce support to pin VMIDs

[Shameer Kolothum]

Introduce kvm_arm_pinned_vmid_get() and kvm_arm_pinned_vmid_put(), to pin
a VMID associated with a KVM instance. This will guarantee that VMID
remains the same after a rollover.

This is in preparation of introducing support in the SMMUv3 driver to use
the KVM VMID for S2 stage configuration in nested mode.

Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
---
 arch/arm64/include/asm/kvm_host.h |  3 ++
 arch/arm64/kvm/vmid.c             | 76 ++++++++++++++++++++++++++++++-
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index d919557af5e5..b6682f5d1b86 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -142,6 +142,7 @@ int topup_hyp_memcache(struct kvm_hyp_memcache *mc, unsigned long min_pages);
 
 struct kvm_vmid {
 	atomic64_t id;
+	refcount_t pinned;
 };
 
 struct kvm_s2_mmu {
@@ -1261,6 +1262,8 @@ int __init kvm_arm_vmid_alloc_init(void);
 void __init kvm_arm_vmid_alloc_free(void);
 void kvm_arm_vmid_update(struct kvm_vmid *kvm_vmid);
 void kvm_arm_vmid_clear_active(void);
+int kvm_arm_pinned_vmid_get(struct kvm *kvm);
+void kvm_arm_pinned_vmid_put(struct kvm *kvm);
 
 static inline void kvm_arm_pvtime_vcpu_init(struct kvm_vcpu_arch *vcpu_arch)
 {
diff --git a/arch/arm64/kvm/vmid.c b/arch/arm64/kvm/vmid.c
index 7fe8ba1a2851..7bda189e927c 100644
--- a/arch/arm64/kvm/vmid.c
+++ b/arch/arm64/kvm/vmid.c
@@ -25,6 +25,10 @@ static unsigned long *vmid_map;
 static DEFINE_PER_CPU(atomic64_t, active_vmids);
 static DEFINE_PER_CPU(u64, reserved_vmids);
 
+static unsigned long max_pinned_vmids;
+static unsigned long nr_pinned_vmids;
+static unsigned long *pinned_vmid_map;
+
 #define VMID_MASK		(~GENMASK(kvm_arm_vmid_bits - 1, 0))
 #define VMID_FIRST_VERSION	(1UL << kvm_arm_vmid_bits)
 
@@ -47,7 +51,10 @@ static void flush_context(void)
 	int cpu;
 	u64 vmid;
 
-	bitmap_zero(vmid_map, NUM_USER_VMIDS);
+	if (pinned_vmid_map)
+		bitmap_copy(vmid_map, pinned_vmid_map, NUM_USER_VMIDS);
+	else
+		bitmap_zero(vmid_map, NUM_USER_VMIDS);
 
 	for_each_possible_cpu(cpu) {
 		vmid = atomic64_xchg_relaxed(&per_cpu(active_vmids, cpu), 0);
@@ -103,6 +110,14 @@ static u64 new_vmid(struct kvm_vmid *kvm_vmid)
 			return newvmid;
 		}
 
+		/*
+		 * If it is pinned, we can keep using it. Note that reserved
+		 * takes priority, because even if it is also pinned, we need to
+		 * update the generation into the reserved_vmids.
+		 */
+		if (refcount_read(&kvm_vmid->pinned))
+			return newvmid;
+
 		if (!__test_and_set_bit(vmid2idx(vmid), vmid_map)) {
 			atomic64_set(&kvm_vmid->id, newvmid);
 			return newvmid;
@@ -169,6 +184,55 @@ void kvm_arm_vmid_update(struct kvm_vmid *kvm_vmid)
 	raw_spin_unlock_irqrestore(&cpu_vmid_lock, flags);
 }
 
+int kvm_arm_pinned_vmid_get(struct kvm *kvm)
+{
+	struct kvm_vmid *kvm_vmid;
+	u64 vmid;
+
+	if (!pinned_vmid_map || !kvm)
+		return -EINVAL;
+
+	kvm_vmid = &kvm->arch.mmu.vmid;
+
+	guard(raw_spinlock_irqsave)(&cpu_vmid_lock);
+	vmid = atomic64_read(&kvm_vmid->id);
+
+	if (refcount_inc_not_zero(&kvm_vmid->pinned))
+		return (vmid & ~VMID_MASK);
+
+	if (nr_pinned_vmids >= max_pinned_vmids)
+		return -EINVAL;
+
+	/*
+	 * If we went through one or more rollover since that VMID was
+	 * used, make sure it is still valid, or generate a new one.
+	 */
+	if (!vmid_gen_match(vmid))
+		vmid = new_vmid(kvm_vmid);
+
+	nr_pinned_vmids++;
+	__set_bit(vmid2idx(vmid), pinned_vmid_map);
+	refcount_set(&kvm_vmid->pinned, 1);
+	return (vmid & ~VMID_MASK);
+}
+
+void kvm_arm_pinned_vmid_put(struct kvm *kvm)
+{
+	struct kvm_vmid *kvm_vmid;
+	u64 vmid;
+
+	if (!pinned_vmid_map || !kvm)
+		return;
+
+	kvm_vmid = &kvm->arch.mmu.vmid;
+	vmid = atomic64_read(&kvm_vmid->id);
+	guard(raw_spinlock_irqsave)(&cpu_vmid_lock);
+	if (refcount_dec_and_test(&kvm_vmid->pinned)) {
+		__clear_bit(vmid2idx(vmid), pinned_vmid_map);
+		nr_pinned_vmids--;
+	}
+}
+
 /*
  * Initialize the VMID allocator
  */
@@ -186,10 +250,20 @@ int __init kvm_arm_vmid_alloc_init(void)
 	if (!vmid_map)
 		return -ENOMEM;
 
+	pinned_vmid_map = bitmap_zalloc(NUM_USER_VMIDS, GFP_KERNEL);
+	nr_pinned_vmids = 0;
+
+	/*
+	 * Ensure we have at least one emty slot available after rollover
+	 * and maximum number of VMIDs are pinned. VMID#0 is reserved.
+	 */
+	max_pinned_vmids = NUM_USER_VMIDS - num_possible_cpus() - 2;
+
 	return 0;
 }
 
 void __init kvm_arm_vmid_alloc_free(void)
 {
+	bitmap_free(pinned_vmid_map);
 	bitmap_free(vmid_map);
 }
-- 
2.47.0

[RFC PATCH v3 2/5] iommufd/device: Associate a kvm pointer to iommufd_device

[Shameer Kolothum]

Add a struct kvm * to iommufd_device_bind() fn and associate it
with idev if bind is successful.

Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
---
 drivers/iommu/iommufd/device.c          | 5 ++++-
 drivers/iommu/iommufd/iommufd_private.h | 2 ++
 drivers/vfio/iommufd.c                  | 2 +-
 include/linux/iommufd.h                 | 4 +++-
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/drivers/iommu/iommufd/device.c b/drivers/iommu/iommufd/device.c
index dfd0898fb6c1..0f3983b16d56 100644
--- a/drivers/iommu/iommufd/device.c
+++ b/drivers/iommu/iommufd/device.c
@@ -146,6 +146,7 @@ void iommufd_device_destroy(struct iommufd_object *obj)
  * iommufd_device_bind - Bind a physical device to an iommu fd
  * @ictx: iommufd file descriptor
  * @dev: Pointer to a physical device struct
+ * @kvm: Pointer to struct kvm if device belongs to a KVM VM
  * @id: Output ID number to return to userspace for this device
  *
  * A successful bind establishes an ownership over the device and returns
@@ -159,7 +160,8 @@ void iommufd_device_destroy(struct iommufd_object *obj)
  * The caller must undo this with iommufd_device_unbind()
  */
 struct iommufd_device *iommufd_device_bind(struct iommufd_ctx *ictx,
-					   struct device *dev, u32 *id)
+					   struct device *dev, struct kvm *kvm,
+					   u32 *id)
 {
 	struct iommufd_device *idev;
 	struct iommufd_group *igroup;
@@ -209,6 +211,7 @@ struct iommufd_device *iommufd_device_bind(struct iommufd_ctx *ictx,
 	if (!iommufd_selftest_is_mock_dev(dev))
 		iommufd_ctx_get(ictx);
 	idev->dev = dev;
+	idev->kvm = kvm;
 	idev->enforce_cache_coherency =
 		device_iommu_capable(dev, IOMMU_CAP_ENFORCE_CACHE_COHERENCY);
 	/* The calling driver is a user until iommufd_device_unbind() */
diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index 0b1bafc7fd99..73201ff2c40e 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -398,6 +398,8 @@ struct iommufd_device {
 	struct list_head group_item;
 	/* always the physical device */
 	struct device *dev;
+	/* ..and kvm if available */
+	struct kvm *kvm;
 	bool enforce_cache_coherency;
 	/* protect iopf_enabled counter */
 	struct mutex iopf_lock;
diff --git a/drivers/vfio/iommufd.c b/drivers/vfio/iommufd.c
index 516294fd901b..664e3579ce0e 100644
--- a/drivers/vfio/iommufd.c
+++ b/drivers/vfio/iommufd.c
@@ -115,7 +115,7 @@ int vfio_iommufd_physical_bind(struct vfio_device *vdev,
 {
 	struct iommufd_device *idev;
 
-	idev = iommufd_device_bind(ictx, vdev->dev, out_device_id);
+	idev = iommufd_device_bind(ictx, vdev->dev, vdev->kvm, out_device_id);
 	if (IS_ERR(idev))
 		return PTR_ERR(idev);
 	vdev->iommufd_device = idev;
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index 11110c749200..ac8cca0190f4 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -22,6 +22,7 @@ struct iommufd_ctx;
 struct iommufd_device;
 struct iommufd_viommu_ops;
 struct page;
+struct kvm;
 
 enum iommufd_object_type {
 	IOMMUFD_OBJ_NONE,
@@ -49,7 +50,8 @@ struct iommufd_object {
 };
 
 struct iommufd_device *iommufd_device_bind(struct iommufd_ctx *ictx,
-					   struct device *dev, u32 *id);
+					   struct device *dev, struct kvm *kvm,
+					   u32 *id);
 void iommufd_device_unbind(struct iommufd_device *idev);
 
 int iommufd_device_attach(struct iommufd_device *idev, u32 *pt_id);
-- 
2.47.0

Re: [RFC PATCH v3 2/5] iommufd/device: Associate a kvm pointer to iommufd_device

[Jason Gunthorpe]

On Wed, Mar 19, 2025 at 05:31:59PM +0000, Shameer Kolothum wrote:
> Add a struct kvm * to iommufd_device_bind() fn and associate it
> with idev if bind is successful.
> 
> Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
> ---
>  drivers/iommu/iommufd/device.c          | 5 ++++-
>  drivers/iommu/iommufd/iommufd_private.h | 2 ++
>  drivers/vfio/iommufd.c                  | 2 +-
>  include/linux/iommufd.h                 | 4 +++-
>  4 files changed, 10 insertions(+), 3 deletions(-)

Reviewed-by: Jason Gunthorpe <jgg@nvidia.com>

I'd probably explain more in the commit message that this is to allow
VFIO to pass the KVM to iommufd so it can later be associated with
downstream objects like the viommu.

Jason

[RFC PATCH v3 3/5] iommu/arm-smmu-v3-iommufd: Pass in kvm pointer to viommu_alloc

[Shameer Kolothum]

No functional changes.

This will be used in a later patch to add support to use
KVM VMID in ARM SMMUv3 s2 stage configuration.

Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c | 1 +
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h         | 1 +
 drivers/iommu/iommufd/viommu.c                      | 3 ++-
 include/linux/iommu.h                               | 4 +++-
 4 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
index 6f8be1164167..ee2fac5c899b 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
@@ -370,6 +370,7 @@ static const struct iommufd_viommu_ops arm_vsmmu_ops = {
 };
 
 struct iommufd_viommu *arm_vsmmu_alloc(struct device *dev,
+				       struct kvm *kvm,
 				       struct iommu_domain *parent,
 				       struct iommufd_ctx *ictx,
 				       unsigned int viommu_type)
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index 1f6696bc4f6c..9f49de52a700 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1060,6 +1060,7 @@ struct arm_vsmmu {
 #if IS_ENABLED(CONFIG_ARM_SMMU_V3_IOMMUFD)
 void *arm_smmu_hw_info(struct device *dev, u32 *length, u32 *type);
 struct iommufd_viommu *arm_vsmmu_alloc(struct device *dev,
+				       struct kvm *kvm,
 				       struct iommu_domain *parent,
 				       struct iommufd_ctx *ictx,
 				       unsigned int viommu_type);
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 69b88e8c7c26..e157d786f295 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -47,7 +47,8 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 		goto out_put_hwpt;
 	}
 
-	viommu = ops->viommu_alloc(idev->dev, hwpt_paging->common.domain,
+	viommu = ops->viommu_alloc(idev->dev, idev->kvm,
+				   hwpt_paging->common.domain,
 				   ucmd->ictx, cmd->type);
 	if (IS_ERR(viommu)) {
 		rc = PTR_ERR(viommu);
diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index cb01fe49f5df..2f61e0178cfa 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -44,6 +44,7 @@ struct iommu_dma_cookie;
 struct iommu_fault_param;
 struct iommufd_ctx;
 struct iommufd_viommu;
+struct kvm;
 
 #define IOMMU_FAULT_PERM_READ	(1 << 0) /* read */
 #define IOMMU_FAULT_PERM_WRITE	(1 << 1) /* write */
@@ -646,7 +647,8 @@ struct iommu_ops {
 	int (*def_domain_type)(struct device *dev);
 
 	struct iommufd_viommu *(*viommu_alloc)(
-		struct device *dev, struct iommu_domain *parent_domain,
+		struct device *dev, struct kvm *kvm,
+		struct iommu_domain *parent_domain,
 		struct iommufd_ctx *ictx, unsigned int viommu_type);
 
 	const struct iommu_domain_ops *default_domain_ops;
-- 
2.47.0

Re: [RFC PATCH v3 3/5] iommu/arm-smmu-v3-iommufd: Pass in kvm pointer to viommu_alloc

[Jason Gunthorpe]

On Wed, Mar 19, 2025 at 05:32:00PM +0000, Shameer Kolothum wrote:
> diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
> index 69b88e8c7c26..e157d786f295 100644
> --- a/drivers/iommu/iommufd/viommu.c
> +++ b/drivers/iommu/iommufd/viommu.c
> @@ -47,7 +47,8 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
>  		goto out_put_hwpt;
>  	}
>  
> -	viommu = ops->viommu_alloc(idev->dev, hwpt_paging->common.domain,
> +	viommu = ops->viommu_alloc(idev->dev, idev->kvm,
> +				   hwpt_paging->common.domain,
>  				   ucmd->ictx, cmd->type);
>  	if (IS_ERR(viommu)) {
>  		rc = PTR_ERR(viommu);

This has a lifetime issue on the kvm pointer.

Because nothing is taking a refcount on the kvm we are relying on the
caller to hold the kvm refcount for the lifetime of the
iommufd_device_bind()/unbind() which is creating the idev.

However, the lifetime of the viommu object is not linked to the
lifetime of the idev. So the idev could be destroyed, and the kvm
refcount put before the viommu is destroyed.

Probably the right answer is to take a refcount on the kvm for the
viommu object somewhere along this path.

Jason

[RFC PATCH v3 4/5] iommu/arm-smmu-v3-iommufd: Use KVM VMID for s2 stage

[Shameer Kolothum]

If kvm is available make use of kvm pinned VMID on BTM
enabled systems to set the s2 stage VMID for nested
domains.

Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
---
 .../arm/arm-smmu-v3/arm-smmu-v3-iommufd.c     | 53 +++++++++++++++++--
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h   |  1 +
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
index ee2fac5c899b..79fcb903741f 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
@@ -3,6 +3,7 @@
  * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES
  */
 
+#include <linux/kvm_host.h>
 #include <uapi/linux/iommufd.h>
 
 #include "arm-smmu-v3.h"
@@ -39,6 +40,48 @@ arm_smmu_get_msi_mapping_domain(struct iommu_domain *domain)
 	return &nested_domain->vsmmu->s2_parent->domain;
 }
 
+static int arm_vsmmu_alloc_vmid(struct arm_smmu_device *smmu, struct kvm *kvm,
+				bool *kvm_used)
+{
+#ifdef CONFIG_KVM
+	/*
+	 * There can only be one allocator for VMIDs active at once. If BTM is
+	 * turned on then KVM's allocator always supplies the VMID, and the
+	 * VMID is matched by CPU invalidation of the KVM S2. Right now there
+	 * is no API to get an unused VMID from KVM so this also means BTM systems
+	 * cannot support S2 without an associated KVM.
+	 */
+	if ((smmu->features & ARM_SMMU_FEAT_BTM)) {
+		int vmid;
+
+		if (!kvm || !kvm_get_kvm_safe(kvm))
+			return -EOPNOTSUPP;
+		vmid = kvm_arm_pinned_vmid_get(kvm);
+		if (vmid < 0)
+			kvm_put_kvm(kvm);
+		else
+			*kvm_used = true;
+		return vmid;
+	}
+#endif
+	return ida_alloc_range(&smmu->vmid_map, 1, (1 << smmu->vmid_bits) - 1,
+			       GFP_KERNEL);
+}
+
+static void arm_vsmmu_free_vmid(struct arm_smmu_device *smmu, u16 vmid,
+				struct kvm *kvm)
+{
+#ifdef CONFIG_KVM
+	if ((smmu->features & ARM_SMMU_FEAT_BTM)) {
+		if (kvm) {
+			kvm_arm_pinned_vmid_put(kvm);
+			return kvm_put_kvm(kvm);
+		}
+	}
+#endif
+	ida_free(&smmu->vmid_map, vmid);
+}
+
 static void arm_vsmmu_destroy(struct iommufd_viommu *viommu)
 {
 	struct arm_vsmmu *vsmmu = container_of(viommu, struct arm_vsmmu, core);
@@ -53,7 +96,7 @@ static void arm_vsmmu_destroy(struct iommufd_viommu *viommu)
 	list_del(&vsmmu->vsmmus_elm);
 	spin_unlock_irqrestore(&vsmmu->s2_parent->vsmmus.lock, flags);
 	arm_smmu_cmdq_issue_cmd_with_sync(smmu, &cmd);
-	ida_free(&smmu->vmid_map, vsmmu->vmid);
+	arm_vsmmu_free_vmid(smmu, vsmmu->vmid, vsmmu->kvm);
 }
 
 static void arm_smmu_make_nested_cd_table_ste(
@@ -379,6 +422,7 @@ struct iommufd_viommu *arm_vsmmu_alloc(struct device *dev,
 		iommu_get_iommu_dev(dev, struct arm_smmu_device, iommu);
 	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
 	struct arm_smmu_domain *s2_parent = to_smmu_domain(parent);
+	bool kvm_used = false;
 	struct arm_vsmmu *vsmmu;
 	unsigned long flags;
 	int vmid;
@@ -409,21 +453,22 @@ struct iommufd_viommu *arm_vsmmu_alloc(struct device *dev,
 	    !(smmu->features & ARM_SMMU_FEAT_S2FWB))
 		return ERR_PTR(-EOPNOTSUPP);
 
-	vmid = ida_alloc_range(&smmu->vmid_map, 1, (1 << smmu->vmid_bits) - 1,
-			       GFP_KERNEL);
+	vmid = arm_vsmmu_alloc_vmid(smmu, kvm, &kvm_used);
 	if (vmid < 0)
 		return ERR_PTR(vmid);
 
 	vsmmu = iommufd_viommu_alloc(ictx, struct arm_vsmmu, core,
 				     &arm_vsmmu_ops);
 	if (IS_ERR(vsmmu)) {
-		ida_free(&smmu->vmid_map, vmid);
+		arm_vsmmu_free_vmid(smmu, vmid, kvm);
 		return ERR_CAST(vsmmu);
 	}
 
 	vsmmu->smmu = smmu;
 	vsmmu->vmid = (u16)vmid;
 	vsmmu->s2_parent = s2_parent;
+	if (kvm_used)
+		vsmmu->kvm = kvm;
 	spin_lock_irqsave(&s2_parent->vsmmus.lock, flags);
 	list_add_tail(&vsmmu->vsmmus_elm, &s2_parent->vsmmus.list);
 	spin_unlock_irqrestore(&s2_parent->vsmmus.lock, flags);
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index 9f49de52a700..5890c233f73b 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1053,6 +1053,7 @@ struct arm_vsmmu {
 	struct arm_smmu_device *smmu;
 	struct arm_smmu_domain *s2_parent;
 	u16 vmid;
+	struct kvm *kvm;
 
 	struct list_head vsmmus_elm; /* arm_smmu_domain::vsmmus::list */
 };
-- 
2.47.0

Re: [RFC PATCH v3 4/5] iommu/arm-smmu-v3-iommufd: Use KVM VMID for s2 stage

[Jason Gunthorpe]

On Wed, Mar 19, 2025 at 05:32:01PM +0000, Shameer Kolothum wrote:
> +static int arm_vsmmu_alloc_vmid(struct arm_smmu_device *smmu, struct kvm *kvm,
> +				bool *kvm_used)
> +{
> +#ifdef CONFIG_KVM
> +	/*
> +	 * There can only be one allocator for VMIDs active at once. If BTM is
> +	 * turned on then KVM's allocator always supplies the VMID, and the
> +	 * VMID is matched by CPU invalidation of the KVM S2. Right now there
> +	 * is no API to get an unused VMID from KVM so this also means BTM systems
> +	 * cannot support S2 without an associated KVM.
> +	 */
> +	if ((smmu->features & ARM_SMMU_FEAT_BTM)) {
> +		int vmid;
> +
> +		if (!kvm || !kvm_get_kvm_safe(kvm))
> +			return -EOPNOTSUPP;

Isn't kvm_get_kvm_safe() in modular KVM code most of the time? You can't
call it like this right?

Per my prior comments, I'm a little nervous to make drivers keep track
of this refcount instead of the core code but this does seem like it
could work..

I also think you should block using the S2 without a viommu & KVM in
BTM mode.. Maybe like this:

@@ -2938,8 +2938,16 @@ static int arm_smmu_attach_dev(struct iommu_domain *domain, struct device *dev)
                cdptr = arm_smmu_alloc_cd_ptr(master, IOMMU_NO_PASID);
                if (!cdptr)
                        return -ENOMEM;
-       } else if (arm_smmu_ssids_in_use(&master->cd_table))
-               return -EBUSY;
+       } else {
+               if (arm_smmu_ssids_in_use(&master->cd_table))
+                       return -EBUSY;
+               /*
+                * S2 cannot be used when BTM is turned on without a VIOMMU and
+                * VMID shared with KVM
+                */
+               if (smmu->features & ARM_SMMU_FEAT_BTM)
+                       return -EOPNOTSUPP;
+       }

That more closely matches how this will eventually work when the VMID
allocation is properly made to be local to the viommu.

Jason

RE: [RFC PATCH v3 4/5] iommu/arm-smmu-v3-iommufd: Use KVM VMID for s2 stage

[Shameerali Kolothum Thodi]

> -----Original Message-----
> From: Jason Gunthorpe <jgg@ziepe.ca>
> Sent: Wednesday, March 19, 2025 11:40 PM
> To: Shameerali Kolothum Thodi <shameerali.kolothum.thodi@huawei.com>
> Cc: kvmarm@lists.linux.dev; iommu@lists.linux.dev; linux-arm-
> kernel@lists.infradead.org; Linuxarm <linuxarm@huawei.com>;
> kevin.tian@intel.com; alex.williamson@redhat.com; maz@kernel.org;
> oliver.upton@linux.dev; will@kernel.org; robin.murphy@arm.com;
> nicolinc@nvidia.com; jean-philippe@linaro.org; Jonathan Cameron
> <jonathan.cameron@huawei.com>
> Subject: Re: [RFC PATCH v3 4/5] iommu/arm-smmu-v3-iommufd: Use KVM
> VMID for s2 stage
> 
> On Wed, Mar 19, 2025 at 05:32:01PM +0000, Shameer Kolothum wrote:
> > +static int arm_vsmmu_alloc_vmid(struct arm_smmu_device *smmu,
> struct kvm *kvm,
> > +				bool *kvm_used)
> > +{
> > +#ifdef CONFIG_KVM
> > +	/*
> > +	 * There can only be one allocator for VMIDs active at once. If BTM
> is
> > +	 * turned on then KVM's allocator always supplies the VMID, and
> the
> > +	 * VMID is matched by CPU invalidation of the KVM S2. Right now
> there
> > +	 * is no API to get an unused VMID from KVM so this also means
> BTM systems
> > +	 * cannot support S2 without an associated KVM.
> > +	 */
> > +	if ((smmu->features & ARM_SMMU_FEAT_BTM)) {
> > +		int vmid;
> > +
> > +		if (!kvm || !kvm_get_kvm_safe(kvm))
> > +			return -EOPNOTSUPP;
> 
> Isn't kvm_get_kvm_safe() in modular KVM code most of the time? You can't
> call it like this right?

I think it is safe since KVM is not modular for ARM64. And that is one of the
reasons I thought I will hold the KVM reference here as we don't have to go
through the symbol_get/put calls here. Also IOMMUFD is not using the kvm
pointer at the moment. Probably I should add some comments in IOMMUFD
to make it clear.

However I just realized that in patch #1 I should EXPORT the
kvm_arm_pinned_vmid_get/put functions as smmuv3 driver can be built
modular. 

> 
> Per my prior comments, I'm a little nervous to make drivers keep track
> of this refcount instead of the core code but this does seem like it
> could work..
> 
> I also think you should block using the S2 without a viommu & KVM in
> BTM mode.. Maybe like this:
> 
> @@ -2938,8 +2938,16 @@ static int arm_smmu_attach_dev(struct
> iommu_domain *domain, struct device *dev)
>                 cdptr = arm_smmu_alloc_cd_ptr(master, IOMMU_NO_PASID);
>                 if (!cdptr)
>                         return -ENOMEM;
> -       } else if (arm_smmu_ssids_in_use(&master->cd_table))
> -               return -EBUSY;
> +       } else {
> +               if (arm_smmu_ssids_in_use(&master->cd_table))
> +                       return -EBUSY;
> +               /*
> +                * S2 cannot be used when BTM is turned on without a VIOMMU
> and
> +                * VMID shared with KVM
> +                */
> +               if (smmu->features & ARM_SMMU_FEAT_BTM)
> +                       return -EOPNOTSUPP;
> +       }
> 
> That more closely matches how this will eventually work when the VMID
> allocation is properly made to be local to the viommu.

Ok. Will update.

Thanks,
Shameer

Re: [RFC PATCH v3 4/5] iommu/arm-smmu-v3-iommufd: Use KVM VMID for s2 stage

[Jason Gunthorpe]

On Thu, Mar 20, 2025 at 09:30:45AM +0000, Shameerali Kolothum Thodi wrote:
> > Isn't kvm_get_kvm_safe() in modular KVM code most of the time? You can't
> > call it like this right?
> 
> I think it is safe since KVM is not modular for ARM64. And that is one of the
> reasons I thought I will hold the KVM reference here as we don't have to go
> through the symbol_get/put calls here. Also IOMMUFD is not using the kvm
> pointer at the moment. Probably I should add some comments in IOMMUFD
> to make it clear.

Oh I didn't know that.. Interesting simplification

> > That more closely matches how this will eventually work when the VMID
> > allocation is properly made to be local to the viommu.
> 
> Ok. Will update.

You should also block creating a VIOMMU without a KVM for the same
reason.. Really isolate the S2 to only being usable with the KVM
shared VMID.

Jason

[RFC PATCH v3 5/5] iommu/arm-smmu-v3: Enable broadcast TLB maintenance

[Shameer Kolothum]

From: Jean-Philippe Brucker <jean-philippe@linaro.org>

The SMMUv3 can handle invalidation targeted at TLB entries with shared
ASIDs. If the implementation supports broadcast TLB maintenance, enable it
and keep track of it in a feature bit. The SMMU will then be affected by
inner-shareable TLB invalidations from other agents.

In order to avoid over invalidation with stage-2 translation contexts,
enable BTM only when SMMUv3 supports eiher S1 or both S1 & S2 transaltion
contexts. In this way the default domain will use stage-1 and stage-2 will
be only used for NESTED Domain setup.

Signed-off-by: Jean-Philippe Brucker <jean-philippe@linaro.org>
[Shameer: Enable BTM only if S1 is supported]
Signed-off-by: Shameer Kolothum <shameerali.kolothum.thodi@huawei.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 24 +++++++++++++++++++--
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h |  1 +
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
index addc6308742b..06a13d78286a 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -4119,11 +4119,14 @@ static int arm_smmu_device_reset(struct arm_smmu_device *smmu)
 	writel_relaxed(reg, smmu->base + ARM_SMMU_CR1);
 
 	/* CR2 (random crap) */
-	reg = CR2_PTM | CR2_RECINVSID;
+	reg = CR2_RECINVSID;
 
 	if (smmu->features & ARM_SMMU_FEAT_E2H)
 		reg |= CR2_E2H;
 
+	if (!(smmu->features & ARM_SMMU_FEAT_BTM))
+		reg |= CR2_PTM;
+
 	writel_relaxed(reg, smmu->base + ARM_SMMU_CR2);
 
 	/* Stream table */
@@ -4289,6 +4292,7 @@ static int arm_smmu_device_hw_probe(struct arm_smmu_device *smmu)
 {
 	u32 reg;
 	bool coherent = smmu->features & ARM_SMMU_FEAT_COHERENCY;
+	bool vhe = cpus_have_cap(ARM64_HAS_VIRT_HOST_EXTN);
 
 	/* IDR0 */
 	reg = readl_relaxed(smmu->base + ARM_SMMU_IDR0);
@@ -4341,7 +4345,7 @@ static int arm_smmu_device_hw_probe(struct arm_smmu_device *smmu)
 
 	if (reg & IDR0_HYP) {
 		smmu->features |= ARM_SMMU_FEAT_HYP;
-		if (cpus_have_cap(ARM64_HAS_VIRT_HOST_EXTN))
+		if (vhe)
 			smmu->features |= ARM_SMMU_FEAT_E2H;
 	}
 
@@ -4368,6 +4372,22 @@ static int arm_smmu_device_hw_probe(struct arm_smmu_device *smmu)
 
 	if (reg & IDR0_S2P)
 		smmu->features |= ARM_SMMU_FEAT_TRANS_S2;
+	/*
+	 * If S1 is supported, verify that BTM can be enabled.  If S2 is available
+	 * and BTM is enabled, S2 will be used exclusively for nested domains,
+	 * ensuring a KVM VMID is obtained.
+	 * BTM is beneficial when the CPU shares page tables with SMMUv3 (e.g., vSVA).
+	 */
+	if (reg & IDR0_S1P) {
+		/*
+		 * If the CPU is using VHE, but the SMMU doesn't support it, the SMMU
+		 * will create TLB entries for NH-EL1 world and will miss the
+		 * broadcasted TLB invalidations that target EL2-E2H world. Don't enable
+		 * BTM in that case.
+		 */
+		if (reg & IDR0_BTM && (!vhe || reg & IDR0_HYP))
+			smmu->features |= ARM_SMMU_FEAT_BTM;
+	}
 
 	if (!(reg & (IDR0_S1P | IDR0_S2P))) {
 		dev_err(smmu->dev, "no translation support!\n");
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index 5890c233f73b..f554b6aa52c9 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -39,6 +39,7 @@ struct arm_smmu_device;
 #define IDR0_HTTU			GENMASK(7, 6)
 #define IDR0_HTTU_ACCESS		1
 #define IDR0_HTTU_ACCESS_DIRTY		2
+#define IDR0_BTM			(1 << 5)
 #define IDR0_COHACC			(1 << 4)
 #define IDR0_TTF			GENMASK(3, 2)
 #define IDR0_TTF_AARCH64		2
-- 
2.47.0