An iotlb tag (ASID/VMID) will not be used:
1) Before being installed to CD/STE during a device attachment
2) After being removed from CD/STE during a device detachment

Both (1) and (2) exactly align with the lifecyle of the domain->invs. So,
it becomes very nature to use domain->invs to allocate/free an ASID/VMID.

Add a pair of function ops in struct arm_smmu_invs, to manage iotlb tag.

Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h |  6 ++
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 93 +++++++++++++++++++++
 2 files changed, 99 insertions(+)

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 0a5aead300b6..b275673c03ce 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -691,6 +691,9 @@ static inline bool arm_smmu_inv_is_ats(const struct arm_smmu_inv *inv)
  * @rwlock: optional rwlock to fench ATS operations
  * @has_ats: flag if the array contains an INV_TYPE_ATS or INV_TYPE_ATS_FULL
  * @rcu: rcu head for kfree_rcu()
+ * @smmu_domain: owner domain of the array
+ * @alloc_id: a callback to allocate a new iotlb tag
+ * @free_id: a callback to free an iotlb tag when its user counter reaches 0
  * @inv: flexible invalidation array
  *
  * The arm_smmu_invs is an RCU data structure. During a ->attach_dev callback,
@@ -720,6 +723,9 @@ struct arm_smmu_invs {
 	rwlock_t rwlock;
 	bool has_ats;
 	struct rcu_head rcu;
+	struct arm_smmu_domain *smmu_domain;
+	int (*alloc_id)(struct arm_smmu_inv *inv, void *data);
+	void (*free_id)(struct arm_smmu_inv *inv, bool flush);
 	struct arm_smmu_inv inv[] __counted_by(max_invs);
 };
 
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 bf0df16cec45..8a2b7064d29b 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -3117,6 +3117,94 @@ static void arm_smmu_disable_iopf(struct arm_smmu_master *master,
 		iopf_queue_remove_device(master->smmu->evtq.iopf, master->dev);
 }
 
+/*
+ * When an array entry's users count reaches zero, it means the ASID/VMID is no
+ * longer being invalidated by map/unmap and must be cleaned. The rule is that
+ * all ASIDs/VMIDs not in an invalidation array are left cleared in the IOTLB.
+ */
+static void arm_smmu_inv_free_asid(struct arm_smmu_inv *inv, bool flush)
+{
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	if (inv->type != INV_TYPE_S1_ASID)
+		return;
+	if (refcount_read(&inv->users))
+		return;
+
+	if (flush) {
+		struct arm_smmu_cmdq_ent cmd = {
+			.opcode = inv->nsize_opcode,
+			.tlbi.asid = inv->id,
+		};
+
+		arm_smmu_cmdq_issue_cmd_with_sync(inv->smmu, &cmd);
+	}
+
+	/* Lastly, free the ASID as the last user detached */
+	xa_erase(&arm_smmu_asid_xa, inv->id);
+}
+
+static void arm_smmu_inv_free_vmid(struct arm_smmu_inv *inv, bool flush)
+{
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	/* Note S2_VMID using nsize_opcode covers S2_VMID_S1_CLEAR already */
+	if (inv->type != INV_TYPE_S2_VMID)
+		return;
+	if (refcount_read(&inv->users))
+		return;
+
+	if (flush) {
+		struct arm_smmu_cmdq_ent cmd = {
+			.opcode = inv->nsize_opcode,
+			.tlbi.vmid = inv->id,
+		};
+
+		arm_smmu_cmdq_issue_cmd_with_sync(inv->smmu, &cmd);
+	}
+
+	/* Lastly, free the VMID as the last user detached */
+	ida_free(&inv->smmu->vmid_map, inv->id);
+}
+
+static int arm_smmu_inv_alloc_asid(struct arm_smmu_inv *inv, void *data)
+{
+	struct arm_smmu_domain *smmu_domain = data;
+	struct arm_smmu_device *smmu = inv->smmu;
+	u32 asid;
+	int ret;
+
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	/* Allocate a new iotlb_tag.id */
+	WARN_ON(inv->type != INV_TYPE_S1_ASID);
+
+	ret = xa_alloc(&arm_smmu_asid_xa, &asid, smmu_domain,
+		       XA_LIMIT(1, (1 << smmu->asid_bits) - 1), GFP_KERNEL);
+	if (ret)
+		return ret;
+	inv->id = asid;
+	return 0;
+}
+
+static int arm_smmu_inv_alloc_vmid(struct arm_smmu_inv *inv, void *data)
+{
+	struct arm_smmu_device *smmu = inv->smmu;
+	int vmid;
+
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	WARN_ON(inv->type != INV_TYPE_S2_VMID);
+
+	/* Reserve VMID 0 for stage-2 bypass STEs */
+	vmid = ida_alloc_range(&smmu->vmid_map, 1, (1 << smmu->vmid_bits) - 1,
+			       GFP_KERNEL);
+	if (vmid < 0)
+		return vmid;
+	inv->id = vmid;
+	return 0;
+}
+
 static struct arm_smmu_inv *
 arm_smmu_master_build_inv(struct arm_smmu_master *master,
 			  enum arm_smmu_inv_type type, u32 id, ioasid_t ssid,
@@ -3191,12 +3279,17 @@ arm_smmu_master_build_invs(struct arm_smmu_master *master, bool ats_enabled,
 					       smmu_domain->cd.asid,
 					       IOMMU_NO_PASID, pgsize))
 			return NULL;
+		master->build_invs->alloc_id = arm_smmu_inv_alloc_asid;
+		master->build_invs->free_id = arm_smmu_inv_free_asid;
+		master->build_invs->smmu_domain = smmu_domain;
 		break;
 	case ARM_SMMU_DOMAIN_S2:
 		if (!arm_smmu_master_build_inv(master, INV_TYPE_S2_VMID,
 					       smmu_domain->s2_cfg.vmid,
 					       IOMMU_NO_PASID, pgsize))
 			return NULL;
+		master->build_invs->alloc_id = arm_smmu_inv_alloc_vmid;
+		master->build_invs->free_id = arm_smmu_inv_free_vmid;
 		break;
 	default:
 		WARN_ON(true);
-- 
2.43.0

On Thu, Dec 18, 2025 at 12:26:48PM -0800, Nicolin Chen wrote:
> @@ -720,6 +723,9 @@ struct arm_smmu_invs {
>  	rwlock_t rwlock;
>  	bool has_ats;
>  	struct rcu_head rcu;
> +	struct arm_smmu_domain *smmu_domain;
> +	int (*alloc_id)(struct arm_smmu_inv *inv, void *data);
> +	void (*free_id)(struct arm_smmu_inv *inv, bool flush);
>  	struct arm_smmu_inv inv[] __counted_by(max_invs);
>  };

I'm not keen on this at all..

On the allocation side everything should be stored inside the
invalidation list, we don't need function callbacks, just get the
information from the list in the first place.

You could probably split it up a bit, introduce arm_smmu_iotlb_tag
first and spread it around, filling it from the domain, then the below
to use the invs array to fill it. But this whole series
should basically just be this:

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 d7c492ee093602..5cd71a9690e292 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -1188,8 +1188,7 @@ EXPORT_SYMBOL_IF_KUNIT(arm_smmu_invs_merge);
  */
 VISIBLE_IF_KUNIT
 void arm_smmu_invs_unref(struct arm_smmu_invs *invs,
-			 struct arm_smmu_invs *to_unref,
-			 void (*free_fn)(struct arm_smmu_inv *inv))
+			 struct arm_smmu_invs *to_unref)
 {
 	unsigned long flags;
 	size_t num_invs = 0;
@@ -1200,16 +1199,16 @@ void arm_smmu_invs_unref(struct arm_smmu_invs *invs,
 		if (cmp < 0) {
 			/* not found in to_unref, leave alone */
 			num_invs = i + 1;
+			refcount_set(&to_unref->inv[j].users, 1);
 		} else if (cmp == 0) {
 			/* same item */
 			if (!refcount_dec_and_test(&invs->inv[i].users)) {
 				num_invs = i + 1;
+				refcount_set(&to_unref->inv[j].users, 1);
 				continue;
 			}
 
-			/* KUNIT test doesn't pass in a free_fn */
-			if (free_fn)
-				free_fn(&invs->inv[i]);
+			refcount_set(&to_unref->inv[j].users, 0);
 			invs->num_trashes++;
 		} else {
 			/* item in to_unref is not in invs or already a trash */
@@ -3046,6 +3045,13 @@ arm_smmu_find_master_domain(struct arm_smmu_domain *smmu_domain,
 	return NULL;
 }
 
+static struct arm_vsmmu *to_vsmmu(struct iommu_domain *domain)
+{
+	if (domain->type == IOMMU_DOMAIN_NESTED)
+		return to_smmu_nested_domain(domain)->vsmmu;
+	return NULL;
+}
+
 /*
  * If the domain uses the smmu_domain->devices list return the arm_smmu_domain
  * structure, otherwise NULL. These domains track attached devices so they can
@@ -3156,12 +3162,136 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
 	case INV_TYPE_ATS:
 	case INV_TYPE_ATS_FULL:
 		cur->size_opcode = cur->nsize_opcode = CMDQ_OP_ATC_INV;
+		cur->ssid = ssid;
 		break;
 	}
 
 	return cur;
 }
 
+static int arm_smmu_get_id_from_invs(struct arm_smmu_domain *smmu_domain,
+				     struct arm_smmu_device *smmu,
+				     unsigned int type)
+{
+	struct arm_smmu_invs *invs = rcu_dereference_protected(
+		smmu_domain->invs, lockdep_is_held(&arm_smmu_asid_lock));
+	size_t i;
+
+	for (i = 0; i != invs->num_invs; i++) {
+		if (invs->inv[i].type == type &&
+		    invs->inv[i].smmu == smmu &&
+		    refcount_read(&invs->inv[i].users))
+		    return invs->inv[i].id;
+	}
+
+	return -ENOENT;
+}
+
+static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
+			    struct arm_smmu_master *master,
+			    struct arm_vsmmu *vsmmu,
+			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
+{
+	struct arm_smmu_device *smmu = master->smmu;
+	int ret;
+	int id;
+
+	lockdep_assert_held(&arm_smmu_asid_lock);
+
+	switch (smmu_domain->stage) {
+	case ARM_SMMU_DOMAIN_SVA:
+	case ARM_SMMU_DOMAIN_S1: {
+		u32 asid;
+
+		if (WARN_ON(vsmmu))
+			return -EINVAL;
+
+		id = arm_smmu_get_id_from_invs(smmu_domain, smmu,
+					       INV_TYPE_S1_ASID);
+		if (id > 0) {
+			tag->asid = id;
+			return 0;
+		}
+
+		if (no_alloc)
+			return -ENOENT;
+
+		ret = xa_alloc(&arm_smmu_asid_xa, &asid, smmu_domain,
+			       XA_LIMIT(1, (1 << smmu->asid_bits) - 1),
+			       GFP_KERNEL);
+		if (ret)
+			return ret;
+		tag->asid = asid;
+		tag->need_free = 1;
+		return 0;
+	}
+
+	case ARM_SMMU_DOMAIN_S2:
+		if (smmu_domain->nest_parent) {
+			/* FIXME we can support attaching a nest_parent without
+			 * a vsmmu, but to do that we need to fix
+			 * arm_smmu_get_id_from_invs() to never return the vmid
+			 * of a vsmmu. Probably by making a
+			 * INV_TYPE_S2_VMID_VSMMU */
+			id = vsmmu->vmid;
+			return 0;
+		}
+
+		id = arm_smmu_get_id_from_invs(smmu_domain, smmu,
+					       INV_TYPE_S2_VMID);
+		if (id > 0) {
+			tag->vmid = id;
+			break;
+		}
+
+		if (no_alloc)
+			return -ENOENT;
+
+		/* Reserve VMID 0 for stage-2 bypass STEs */
+		id = ida_alloc_range(&smmu->vmid_map, 1,
+				     (1 << smmu->vmid_bits) - 1, GFP_KERNEL);
+		if (id < 0)
+			return id;
+		tag->vmid = id;
+		tag->need_free = 1;
+		return 0;
+	}
+	return -EINVAL;
+}
+
+static void arm_smmu_free_tag(struct arm_smmu_master *master,
+			      struct arm_smmu_domain *smmu_domain,
+			      struct arm_smmu_iotlb_tag *tag,
+			      struct arm_smmu_inv *inv)
+{
+	struct arm_smmu_device *smmu = master->smmu;
+	struct arm_smmu_cmdq_ent cmd = {};
+
+	switch (smmu_domain->stage) {
+	case ARM_SMMU_DOMAIN_SVA:
+	case ARM_SMMU_DOMAIN_S1:
+		cmd.opcode = inv->nsize_opcode;
+		cmd.tlbi.asid = tag->asid;
+		arm_smmu_cmdq_issue_cmd_with_sync(smmu, &cmd);
+		xa_erase(&arm_smmu_asid_xa, tag->asid);
+		return;
+
+	case ARM_SMMU_DOMAIN_S2:
+		cmd.opcode = inv->nsize_opcode;
+		cmd.tlbi.vmid = tag->vmid;
+		arm_smmu_cmdq_issue_cmd_with_sync(smmu, &cmd);
+
+		/*
+		 * Even though the VMID remains with the viommu and is not freed
+		 * we are stopping invalidations of it so it has to be cleared.
+		 */
+		if (smmu_domain->nest_parent)
+			return;
+		ida_free(&smmu->vmid_map, tag->vmid);
+		return;
+	}
+}
+
 /*
  * Use the preallocated scratch array at master->build_invs, to build a to_merge
  * or to_unref array, to pass into a following arm_smmu_invs_merge/unref() call.
@@ -3171,7 +3301,8 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
  */
 static struct arm_smmu_invs *
 arm_smmu_master_build_invs(struct arm_smmu_master *master, bool ats_enabled,
-			   ioasid_t ssid, struct arm_smmu_domain *smmu_domain)
+			   ioasid_t ssid, struct arm_smmu_domain *smmu_domain,
+			   struct arm_smmu_iotlb_tag *tag)
 {
 	const bool nesting = smmu_domain->nest_parent;
 	size_t pgsize = 0, i;
@@ -3188,14 +3319,14 @@ arm_smmu_master_build_invs(struct arm_smmu_master *master, bool ats_enabled,
 	case ARM_SMMU_DOMAIN_SVA:
 	case ARM_SMMU_DOMAIN_S1:
 		if (!arm_smmu_master_build_inv(master, INV_TYPE_S1_ASID,
-					       smmu_domain->cd.asid,
-					       IOMMU_NO_PASID, pgsize))
+					       tag->asid, IOMMU_NO_PASID,
+					       pgsize))
 			return NULL;
 		break;
 	case ARM_SMMU_DOMAIN_S2:
 		if (!arm_smmu_master_build_inv(master, INV_TYPE_S2_VMID,
-					       smmu_domain->s2_cfg.vmid,
-					       IOMMU_NO_PASID, pgsize))
+					       tag->vmid, IOMMU_NO_PASID,
+					       pgsize))
 			return NULL;
 		break;
 	default:
@@ -3205,9 +3336,9 @@ arm_smmu_master_build_invs(struct arm_smmu_master *master, bool ats_enabled,
 
 	/* All the nested S1 ASIDs have to be flushed when S2 parent changes */
 	if (nesting) {
-		if (!arm_smmu_master_build_inv(
-			    master, INV_TYPE_S2_VMID_S1_CLEAR,
-			    smmu_domain->s2_cfg.vmid, IOMMU_NO_PASID, 0))
+		if (!arm_smmu_master_build_inv(master,
+					       INV_TYPE_S2_VMID_S1_CLEAR,
+					       tag->vmid, IOMMU_NO_PASID, 0))
 			return NULL;
 	}
 
@@ -3270,12 +3401,14 @@ static void arm_smmu_remove_master_domain(struct arm_smmu_master *master,
  * we are sequencing to update them.
  */
 static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
+					struct iommu_domain *new_domain,
 					struct arm_smmu_domain *new_smmu_domain)
 {
 	struct arm_smmu_domain *old_smmu_domain =
 		to_smmu_domain_devices(state->old_domain);
 	struct arm_smmu_master *master = state->master;
 	ioasid_t ssid = state->ssid;
+	int ret;
 
 	/*
 	 * At this point a NULL domain indicates the domain doesn't use the
@@ -3289,8 +3422,18 @@ static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
 		invst->old_invs = rcu_dereference_protected(
 			new_smmu_domain->invs,
 			lockdep_is_held(&arm_smmu_asid_lock));
-		build_invs = arm_smmu_master_build_invs(
-			master, state->ats_enabled, ssid, new_smmu_domain);
+
+		/* Re-use or allocate an IOTLB cache tag */
+		ret = arm_smmu_get_tag(new_smmu_domain, master,
+				       to_vsmmu(new_domain), &invst->tag,
+				       false);
+		if (ret)
+			return ret;
+
+		build_invs = arm_smmu_master_build_invs(master,
+							state->ats_enabled,
+							ssid, new_smmu_domain,
+							&invst->tag);
 		if (!build_invs)
 			return -EINVAL;
 
@@ -3311,9 +3454,18 @@ static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
 			invst->old_invs = rcu_dereference_protected(
 				old_smmu_domain->invs,
 				lockdep_is_held(&arm_smmu_asid_lock));
+
+		/* Find the IOTLB cache tag this master was using */
+		ret = arm_smmu_get_tag(old_smmu_domain, master,
+				       to_vsmmu(state->old_domain), &invst->tag,
+				       true);
+		if (WARN_ON(ret))
+			return ret;
+
 		/* For old_smmu_domain, new_invs points to master->build_invs */
 		invst->new_invs = arm_smmu_master_build_invs(
-			master, master->ats_enabled, ssid, old_smmu_domain);
+			master, master->ats_enabled, ssid, old_smmu_domain,
+			&invst->tag);
 	}
 
 	return 0;
@@ -3349,36 +3501,13 @@ arm_smmu_install_new_domain_invs(struct arm_smmu_attach_state *state)
 	kfree_rcu(invst->old_invs, rcu);
 }
 
-/*
- * When an array entry's users count reaches zero, it means the ASID/VMID is no
- * longer being invalidated by map/unmap and must be cleaned. The rule is that
- * all ASIDs/VMIDs not in an invalidation array are left cleared in the IOTLB.
- */
-static void arm_smmu_inv_flush_iotlb_tag(struct arm_smmu_inv *inv)
-{
-	struct arm_smmu_cmdq_ent cmd = {};
-
-	switch (inv->type) {
-	case INV_TYPE_S1_ASID:
-		cmd.tlbi.asid = inv->id;
-		break;
-	case INV_TYPE_S2_VMID:
-		/* S2_VMID using nsize_opcode covers S2_VMID_S1_CLEAR */
-		cmd.tlbi.vmid = inv->id;
-		break;
-	default:
-		return;
-	}
-
-	cmd.opcode = inv->nsize_opcode;
-	arm_smmu_cmdq_issue_cmd_with_sync(inv->smmu, &cmd);
-}
-
 /* Should be installed after arm_smmu_install_ste_for_dev() */
 static void
 arm_smmu_install_old_domain_invs(struct arm_smmu_attach_state *state)
 {
 	struct arm_smmu_inv_state *invst = &state->old_domain_invst;
+	struct arm_smmu_domain *old_smmu_domain =
+		to_smmu_domain_devices(state->old_domain);
 	struct arm_smmu_invs *old_invs = invst->old_invs;
 	struct arm_smmu_invs *new_invs;
 
@@ -3387,8 +3516,10 @@ arm_smmu_install_old_domain_invs(struct arm_smmu_attach_state *state)
 	if (!invst->invs_ptr)
 		return;
 
-	arm_smmu_invs_unref(old_invs, invst->new_invs,
-			    arm_smmu_inv_flush_iotlb_tag);
+	arm_smmu_invs_unref(old_invs, invst->new_invs);
+	if (old_smmu_domain && !refcount_read(&invst->new_invs->inv[0].users))
+		arm_smmu_free_tag(state->master, old_smmu_domain, &invst->tag,
+				  invst->new_invs->inv);
 
 	new_invs = arm_smmu_invs_purge(old_invs);
 	if (!new_invs)
@@ -3472,9 +3603,9 @@ int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 				     arm_smmu_ats_supported(master);
 	}
 
-	ret = arm_smmu_attach_prepare_invs(state, smmu_domain);
+	ret = arm_smmu_attach_prepare_invs(state, new_domain, smmu_domain);
 	if (ret)
-		return ret;
+		goto err_unprepare_invs;
 
 	if (smmu_domain) {
 		if (new_domain->type == IOMMU_DOMAIN_NESTED) {
@@ -3551,6 +3682,10 @@ int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 err_free_vmaster:
 	kfree(state->vmaster);
 err_unprepare_invs:
+	if (state->new_domain_invst.tag.need_free)
+		arm_smmu_free_tag(master, smmu_domain,
+				  &state->new_domain_invst.tag,
+				  state->new_domain_invst.new_invs.inv);
 	kfree(state->new_domain_invst.new_invs);
 	return ret;
 }
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 4f104c1baa67a2..689b5ca7d61525 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1103,6 +1103,12 @@ static inline bool arm_smmu_master_canwbs(struct arm_smmu_master *master)
 	       IOMMU_FWSPEC_PCI_RC_CANWBS;
 }
 
+struct arm_smmu_iotlb_tag {
+	u16 need_free : 1;
+	u16 asid;
+	u16 vmid;
+};
+
 /**
  * struct arm_smmu_inv_state - Per-domain invalidation array state
  * @invs_ptr: points to the domain->invs (unwinding nesting/etc.) or is NULL if
@@ -1116,6 +1122,7 @@ struct arm_smmu_inv_state {
 	struct arm_smmu_invs __rcu **invs_ptr;
 	struct arm_smmu_invs *old_invs;
 	struct arm_smmu_invs *new_invs;
+	struct arm_smmu_iotlb_tag tag;
 };
 
 struct arm_smmu_attach_state {

Hi Jason,

On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> +			    struct arm_smmu_master *master,
> +			    struct arm_vsmmu *vsmmu,
> +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
[...]
> +	case ARM_SMMU_DOMAIN_S2:
> +		if (smmu_domain->nest_parent) {
> +			/* FIXME we can support attaching a nest_parent without
> +			 * a vsmmu, but to do that we need to fix
> +			 * arm_smmu_get_id_from_invs() to never return the vmid
> +			 * of a vsmmu. Probably by making a
> +			 * INV_TYPE_S2_VMID_VSMMU */
> +			id = vsmmu->vmid;
> +			return 0;
> +		}

Would you mind elaborating why arm_smmu_get_id_from_invs() can't
return vsmmu->vmid to share with a naked S2 STE?

I'm having a bit trouble justifying this INV_TYPE_S2_VMID_VSMMU.

Since it's the same S2 domain/iopt, if anything that is attached
(whether nested or naked) changes the S2 iopt, we should always
flush the nested S1 domain too, right?

If so, a naked S2 STE should refer to the same VMID in order to
allow its following INV_TYPE_S2_VMID_S1_CLEAR to work properly?

Thanks
Nicolin

On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> Hi Jason,
> 
> On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > +			    struct arm_smmu_master *master,
> > +			    struct arm_vsmmu *vsmmu,
> > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> [...]
> > +	case ARM_SMMU_DOMAIN_S2:
> > +		if (smmu_domain->nest_parent) {
> > +			/* FIXME we can support attaching a nest_parent without
> > +			 * a vsmmu, but to do that we need to fix
> > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > +			 * of a vsmmu. Probably by making a
> > +			 * INV_TYPE_S2_VMID_VSMMU */
> > +			id = vsmmu->vmid;
> > +			return 0;
> > +		}
> 
> Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> return vsmmu->vmid to share with a naked S2 STE?

A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
impossible to get vsmmu->vmid.

The only domains which have it are nested domains using a bypass vSTE.

So, if userspace attaches a bypass vSTE then it should use vsmmu->vmid

However if they attach a raw S2 HWPT without a vSTE then there is no
vsmmu and it should work like any other S2 attach and allocate a VMID
for this domain, ignoring any vSMMU that may exist.

In sort, the only case where we use the vsmmu->vmid is for vSTEs.

Jason

On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > Hi Jason,
> > 
> > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > +			    struct arm_smmu_master *master,
> > > +			    struct arm_vsmmu *vsmmu,
> > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > [...]
> > > +	case ARM_SMMU_DOMAIN_S2:
> > > +		if (smmu_domain->nest_parent) {
> > > +			/* FIXME we can support attaching a nest_parent without
> > > +			 * a vsmmu, but to do that we need to fix
> > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > +			 * of a vsmmu. Probably by making a
> > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > +			id = vsmmu->vmid;
> > > +			return 0;
> > > +		}
> > 
> > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > return vsmmu->vmid to share with a naked S2 STE?
> 
> A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> impossible to get vsmmu->vmid.

An S2 parent domain should be per VM. And a vSMMU on top of an S2
should be per SMMU. So, it could have stored a list of vSMMUs and
device attaching to a naked S2 could match its master->smmu with
vSMMU->smmu in the list?

But, it's probably not very efficient, so I took your advise and
added an INV_TYPE_S2_VMID_VSMMU. Overall, this looks clean to me.

> The only domains which have it are nested domains using a bypass vSTE.
> 
> So, if userspace attaches a bypass vSTE then it should use vsmmu->vmid
> 
> However if they attach a raw S2 HWPT without a vSTE then there is no
> vsmmu and it should work like any other S2 attach and allocate a VMID
> for this domain, ignoring any vSMMU that may exist.
> 
> In sort, the only case where we use the vsmmu->vmid is for vSTEs.

Yea, I think I got these right. I have all the code reworked and
will run some tests and post v2.

Thanks
Nicolin

On Thu, Jan 15, 2026 at 09:13:00PM -0800, Nicolin Chen wrote:
> On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> > On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > > Hi Jason,
> > > 
> > > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > > +			    struct arm_smmu_master *master,
> > > > +			    struct arm_vsmmu *vsmmu,
> > > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > > [...]
> > > > +	case ARM_SMMU_DOMAIN_S2:
> > > > +		if (smmu_domain->nest_parent) {
> > > > +			/* FIXME we can support attaching a nest_parent without
> > > > +			 * a vsmmu, but to do that we need to fix
> > > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > > +			 * of a vsmmu. Probably by making a
> > > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > > +			id = vsmmu->vmid;
> > > > +			return 0;
> > > > +		}
> > > 
> > > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > > return vsmmu->vmid to share with a naked S2 STE?
> > 
> > A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> > impossible to get vsmmu->vmid.
> 
> An S2 parent domain should be per VM. And a vSMMU on top of an S2
> should be per SMMU. So, it could have stored a list of vSMMUs and
> device attaching to a naked S2 could match its master->smmu with
> vSMMU->smmu in the list?

That would cause lifecycle problems if the vSMMU is destroyed
while the nake S2 is still attached and trying to use the vSMMU's
VMID.

Jason

On Fri, Jan 16, 2026 at 10:41:20AM -0400, Jason Gunthorpe wrote:
> On Thu, Jan 15, 2026 at 09:13:00PM -0800, Nicolin Chen wrote:
> > On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> > > On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > > > Hi Jason,
> > > > 
> > > > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > > > +			    struct arm_smmu_master *master,
> > > > > +			    struct arm_vsmmu *vsmmu,
> > > > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > > > [...]
> > > > > +	case ARM_SMMU_DOMAIN_S2:
> > > > > +		if (smmu_domain->nest_parent) {
> > > > > +			/* FIXME we can support attaching a nest_parent without
> > > > > +			 * a vsmmu, but to do that we need to fix
> > > > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > > > +			 * of a vsmmu. Probably by making a
> > > > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > > > +			id = vsmmu->vmid;
> > > > > +			return 0;
> > > > > +		}
> > > > 
> > > > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > > > return vsmmu->vmid to share with a naked S2 STE?
> > > 
> > > A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> > > impossible to get vsmmu->vmid.
> > 
> > An S2 parent domain should be per VM. And a vSMMU on top of an S2
> > should be per SMMU. So, it could have stored a list of vSMMUs and
> > device attaching to a naked S2 could match its master->smmu with
> > vSMMU->smmu in the list?
> 
> That would cause lifecycle problems if the vSMMU is destroyed
> while the nake S2 is still attached and trying to use the vSMMU's
> VMID.

Well, if vSMMU code does the same get-build-merge in vsmmu_init()
and build-unref in vsmmu_destroy(), VMID is basically managed by
the invalidation array. Yes, a naked S2 attachment would still use
the shared VMID, but I think that's fine for a nesting parent?

That being said, the implementation isn't going to be as clean as
your approach.

Nicolin

On Fri, Jan 16, 2026 at 08:58:02AM -0800, Nicolin Chen wrote:
> On Fri, Jan 16, 2026 at 10:41:20AM -0400, Jason Gunthorpe wrote:
> > On Thu, Jan 15, 2026 at 09:13:00PM -0800, Nicolin Chen wrote:
> > > On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> > > > On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > > > > Hi Jason,
> > > > > 
> > > > > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > > > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > > > > +			    struct arm_smmu_master *master,
> > > > > > +			    struct arm_vsmmu *vsmmu,
> > > > > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > > > > [...]
> > > > > > +	case ARM_SMMU_DOMAIN_S2:
> > > > > > +		if (smmu_domain->nest_parent) {
> > > > > > +			/* FIXME we can support attaching a nest_parent without
> > > > > > +			 * a vsmmu, but to do that we need to fix
> > > > > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > > > > +			 * of a vsmmu. Probably by making a
> > > > > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > > > > +			id = vsmmu->vmid;
> > > > > > +			return 0;
> > > > > > +		}
> > > > > 
> > > > > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > > > > return vsmmu->vmid to share with a naked S2 STE?
> > > > 
> > > > A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> > > > impossible to get vsmmu->vmid.
> > > 
> > > An S2 parent domain should be per VM. And a vSMMU on top of an S2
> > > should be per SMMU. So, it could have stored a list of vSMMUs and
> > > device attaching to a naked S2 could match its master->smmu with
> > > vSMMU->smmu in the list?
> > 
> > That would cause lifecycle problems if the vSMMU is destroyed
> > while the nake S2 is still attached and trying to use the vSMMU's
> > VMID.
> 
> Well, if vSMMU code does the same get-build-merge in vsmmu_init()
> and build-unref in vsmmu_destroy(), VMID is basically managed by
> the invalidation array. Yes, a naked S2 attachment would still use
> the shared VMID, but I think that's fine for a nesting parent?

The VMID has to be managed by the vsmmu itself, it would be a big
confusing mess any other way. There are multiple invalidation lists,
so you can't rely on that to refcount things.

Jason

On Fri, Jan 16, 2026 at 01:11:58PM -0400, Jason Gunthorpe wrote:
> On Fri, Jan 16, 2026 at 08:58:02AM -0800, Nicolin Chen wrote:
> > On Fri, Jan 16, 2026 at 10:41:20AM -0400, Jason Gunthorpe wrote:
> > > On Thu, Jan 15, 2026 at 09:13:00PM -0800, Nicolin Chen wrote:
> > > > On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> > > > > On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > > > > > Hi Jason,
> > > > > > 
> > > > > > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > > > > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > > > > > +			    struct arm_smmu_master *master,
> > > > > > > +			    struct arm_vsmmu *vsmmu,
> > > > > > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > > > > > [...]
> > > > > > > +	case ARM_SMMU_DOMAIN_S2:
> > > > > > > +		if (smmu_domain->nest_parent) {
> > > > > > > +			/* FIXME we can support attaching a nest_parent without
> > > > > > > +			 * a vsmmu, but to do that we need to fix
> > > > > > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > > > > > +			 * of a vsmmu. Probably by making a
> > > > > > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > > > > > +			id = vsmmu->vmid;
> > > > > > > +			return 0;
> > > > > > > +		}
> > > > > > 
> > > > > > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > > > > > return vsmmu->vmid to share with a naked S2 STE?
> > > > > 
> > > > > A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> > > > > impossible to get vsmmu->vmid.
> > > > 
> > > > An S2 parent domain should be per VM. And a vSMMU on top of an S2
> > > > should be per SMMU. So, it could have stored a list of vSMMUs and
> > > > device attaching to a naked S2 could match its master->smmu with
> > > > vSMMU->smmu in the list?
> > > 
> > > That would cause lifecycle problems if the vSMMU is destroyed
> > > while the nake S2 is still attached and trying to use the vSMMU's
> > > VMID.
> > 
> > Well, if vSMMU code does the same get-build-merge in vsmmu_init()
> > and build-unref in vsmmu_destroy(), VMID is basically managed by
> > the invalidation array. Yes, a naked S2 attachment would still use
> > the shared VMID, but I think that's fine for a nesting parent?
> 
> The VMID has to be managed by the vsmmu itself, it would be a big
> confusing mess any other way. There are multiple invalidation lists,
> so you can't rely on that to refcount things.

Hmm, in which case will we have multiple invalidation lists? And,
why can't we rely on the refcount? (for learning purpose here)

My view is that VM has one big beautiful invalidation list on the
S2 parent domain, which manages all the VMIDs. So, basically VMIDs
could be still tied to the S2 parent domain on attachments. And a
VMID must be one per SMMU instance, whether initially allocated by
a naked or a nested attachment. A second user would just increase
the refcount.

This could work for standard vSMMU case, as its VMID would be only
used by the invalidation function that must be after an attachment
to a nested domain.

Yet, the CMDQ-V case made it complicated, because HW spec suggests
to program the VMID when enabling a VTINF..

Thanks
Nicolin

On Fri, Jan 16, 2026 at 10:27:25AM -0800, Nicolin Chen wrote:
> On Fri, Jan 16, 2026 at 01:11:58PM -0400, Jason Gunthorpe wrote:
> > On Fri, Jan 16, 2026 at 08:58:02AM -0800, Nicolin Chen wrote:
> > > On Fri, Jan 16, 2026 at 10:41:20AM -0400, Jason Gunthorpe wrote:
> > > > On Thu, Jan 15, 2026 at 09:13:00PM -0800, Nicolin Chen wrote:
> > > > > On Fri, Jan 02, 2026 at 11:57:15AM -0400, Jason Gunthorpe wrote:
> > > > > > On Tue, Dec 30, 2025 at 10:52:53AM -0800, Nicolin Chen wrote:
> > > > > > > Hi Jason,
> > > > > > > 
> > > > > > > On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > > > > > > > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > > > > > > > +			    struct arm_smmu_master *master,
> > > > > > > > +			    struct arm_vsmmu *vsmmu,
> > > > > > > > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> > > > > > > [...]
> > > > > > > > +	case ARM_SMMU_DOMAIN_S2:
> > > > > > > > +		if (smmu_domain->nest_parent) {
> > > > > > > > +			/* FIXME we can support attaching a nest_parent without
> > > > > > > > +			 * a vsmmu, but to do that we need to fix
> > > > > > > > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > > > > > > > +			 * of a vsmmu. Probably by making a
> > > > > > > > +			 * INV_TYPE_S2_VMID_VSMMU */
> > > > > > > > +			id = vsmmu->vmid;
> > > > > > > > +			return 0;
> > > > > > > > +		}
> > > > > > > 
> > > > > > > Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> > > > > > > return vsmmu->vmid to share with a naked S2 STE?
> > > > > > 
> > > > > > A "naked" S2 domain doesn't have a pointer to the vsmmu, so it is
> > > > > > impossible to get vsmmu->vmid.
> > > > > 
> > > > > An S2 parent domain should be per VM. And a vSMMU on top of an S2
> > > > > should be per SMMU. So, it could have stored a list of vSMMUs and
> > > > > device attaching to a naked S2 could match its master->smmu with
> > > > > vSMMU->smmu in the list?
> > > > 
> > > > That would cause lifecycle problems if the vSMMU is destroyed
> > > > while the nake S2 is still attached and trying to use the vSMMU's
> > > > VMID.
> > > 
> > > Well, if vSMMU code does the same get-build-merge in vsmmu_init()
> > > and build-unref in vsmmu_destroy(), VMID is basically managed by
> > > the invalidation array. Yes, a naked S2 attachment would still use
> > > the shared VMID, but I think that's fine for a nesting parent?
> > 
> > The VMID has to be managed by the vsmmu itself, it would be a big
> > confusing mess any other way. There are multiple invalidation lists,
> > so you can't rely on that to refcount things.
> 
> Hmm, in which case will we have multiple invalidation lists? And,
> why can't we rely on the refcount? (for learning purpose here)

Well, you could have multiple S2s floating around the system, you'd
need to have some way to figure out which S2 is which nesting parent
for which viommu.

> My view is that VM has one big beautiful invalidation list on the
> S2 parent domain, which manages all the VMIDs. So, basically VMIDs
> could be still tied to the S2 parent domain on attachments. And a
> VMID must be one per SMMU instance, whether initially allocated by
> a naked or a nested attachment. A second user would just increase
> the refcount.

I guess that could work but it makes it much more complicated IMHO.

It would be simpler if there is only ever one refcount on the VMID for
the vsmmu in the S2's list and that goes away when the VSMMU is
destroyed.

Jason

On Tue, Dec 30, 2025 at 10:52:55AM -0800, Nicolin Chen wrote:
> Hi Jason,
> 
> On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> > +static int arm_smmu_get_tag(struct arm_smmu_domain *smmu_domain,
> > +			    struct arm_smmu_master *master,
> > +			    struct arm_vsmmu *vsmmu,
> > +			    struct arm_smmu_iotlb_tag *tag, bool no_alloc)
> [...]
> > +	case ARM_SMMU_DOMAIN_S2:
> > +		if (smmu_domain->nest_parent) {
> > +			/* FIXME we can support attaching a nest_parent without
> > +			 * a vsmmu, but to do that we need to fix
> > +			 * arm_smmu_get_id_from_invs() to never return the vmid
> > +			 * of a vsmmu. Probably by making a
> > +			 * INV_TYPE_S2_VMID_VSMMU */
> > +			id = vsmmu->vmid;
> > +			return 0;
> > +		}
> 
> Would you mind elaborating why arm_smmu_get_id_from_invs() can't
> return vsmmu->vmid to share with a naked S2 STE?
> 
> I'm having a bit trouble justifying this INV_TYPE_S2_VMID_VSMMU.
> 
> Since it's the same S2 domain/iopt, if anything that is attached
> (whether nested or naked) changes the S2 iopt, we should always
> flush the nested S1 domain too, right?
> 
> If so, a naked S2 STE should refer to the same VMID in order to
> allow its following INV_TYPE_S2_VMID_S1_CLEAR to work properly?

I figured that having separate VMIDs on a nest_parent S2 isn't
a problem since invalidation will go through all the VMIDs. And
having an INV_TYPE_S2_VMID_VSMMU is for its unique lifecycle.

Nicolin

On Fri, Dec 19, 2025 at 01:05:51PM -0400, Jason Gunthorpe wrote:
> On Thu, Dec 18, 2025 at 12:26:48PM -0800, Nicolin Chen wrote:
> > @@ -720,6 +723,9 @@ struct arm_smmu_invs {
> >  	rwlock_t rwlock;
> >  	bool has_ats;
> >  	struct rcu_head rcu;
> > +	struct arm_smmu_domain *smmu_domain;
> > +	int (*alloc_id)(struct arm_smmu_inv *inv, void *data);
> > +	void (*free_id)(struct arm_smmu_inv *inv, bool flush);
> >  	struct arm_smmu_inv inv[] __counted_by(max_invs);
> >  };
> 
> I'm not keen on this at all..
> 
> On the allocation side everything should be stored inside the
> invalidation list, we don't need function callbacks, just get the
> information from the list in the first place.
[...]
> @@ -3289,8 +3422,18 @@ static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
>  		invst->old_invs = rcu_dereference_protected(
>  			new_smmu_domain->invs,
>  			lockdep_is_held(&arm_smmu_asid_lock));
> -		build_invs = arm_smmu_master_build_invs(
> -			master, state->ats_enabled, ssid, new_smmu_domain);
> +
> +		/* Re-use or allocate an IOTLB cache tag */
> +		ret = arm_smmu_get_tag(new_smmu_domain, master,
> +				       to_vsmmu(new_domain), &invst->tag,
> +				       false);
> +		if (ret)
> +			return ret;
> +
> +		build_invs = arm_smmu_master_build_invs(master,
> +							state->ats_enabled,
> +							ssid, new_smmu_domain,
> +							&invst->tag);
>  		if (!build_invs)
>  			return -EINVAL;

I had this when I was drafting the series in the beginning, but
later changed to the callback functions :-/

Again, my intention is to avoid iterating the array since merge
and unref functions are already doing that in their first loops.

That being said, if this get-before-build is the preferred way,
let's do that.

Thanks
Nicolin

On Fri, Dec 19, 2025 at 12:56:17PM -0800, Nicolin Chen wrote:
> Again, my intention is to avoid iterating the array since merge
> and unref functions are already doing that in their first loops.

It is not a performance path, and double iterating like this is much
easier to implement with clearer error flows.

Jason