An IOTLB tag now is forwarded from arm_smmu_domain_get_iotlb_tag() to its
final destination (a CD or STE entry).

Thus, arm_smmu_domain_get_iotlb_tag() can safely delink its references to
the cd->asid and s2_cfg->vmid in the smmu_domain. Instead, allocate a new
IOTLB cache tag from the xarray/ida.

The old asid and vmid in the smmu_domain will be deprecated, once VMID is
decoupled from the vSMMU use case.

Suggested-by: Jason Gunthorpe <jgg@nvidia.com>
Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 38 ++++++++++++++++++---
 1 file changed, 33 insertions(+), 5 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 1927eb794db9..d10593823353 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -1210,6 +1210,8 @@ void arm_smmu_invs_unref(struct arm_smmu_invs *invs,
 			/* KUNIT test doesn't pass in a free_fn */
 			if (free_fn)
 				free_fn(&invs->inv[i]);
+			/* Notify the caller to free the iotlb tag */
+			refcount_set(&to_unref->inv[j].users, 0);
 			invs->num_trashes++;
 		} else {
 			/* item in to_unref is not in invs or already a trash */
@@ -3165,12 +3167,31 @@ int arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
 	if (!ret || !alloc)
 		return ret;
 
-	if (tag->type == INV_TYPE_S1_ASID)
-		tag->id = smmu_domain->cd.asid;
-	else
-		tag->id = smmu_domain->s2_cfg.vmid;
+	/* Allocate a new IOTLB cache tag (users counter == 0) */
+	lockdep_assert_held(&arm_smmu_asid_lock);
 
-	return 0;
+	if (tag->type == INV_TYPE_S1_ASID) {
+		ret = xa_alloc(&arm_smmu_asid_xa, &tag->id, smmu_domain,
+			       XA_LIMIT(1, (1 << smmu->asid_bits) - 1),
+			       GFP_KERNEL);
+	} else {
+		ret = ida_alloc_range(&smmu->vmid_map, 1,
+				      (1 << smmu->vmid_bits) - 1, GFP_KERNEL);
+		if (ret > 0) {
+			tag->id = ret; /* int is good for 16-bit VMID */
+			ret = 0;
+		}
+	}
+
+	return ret;
+}
+
+static void arm_smmu_iotlb_tag_free(struct arm_smmu_inv *tag)
+{
+	if (tag->type == INV_TYPE_S1_ASID)
+		xa_erase(&arm_smmu_asid_xa, tag->id);
+	else if (tag->type == INV_TYPE_S2_VMID)
+		ida_free(&tag->smmu->vmid_map, tag->id);
 }
 
 static struct arm_smmu_inv *
@@ -3220,6 +3241,9 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
 		break;
 	}
 
+	/* Set a default users counter */
+	refcount_set(&cur->users, 1);
+
 	return cur;
 }
 
@@ -3453,6 +3477,8 @@ arm_smmu_install_old_domain_invs(struct arm_smmu_attach_state *state)
 
 	arm_smmu_invs_unref(old_invs, invst->new_invs,
 			    arm_smmu_inv_flush_iotlb_tag);
+	if (!refcount_read(&invst->new_invs->inv[0].users))
+		arm_smmu_iotlb_tag_free(&invst->tag);
 
 	new_invs = arm_smmu_invs_purge(old_invs);
 	if (!new_invs)
@@ -3615,6 +3641,8 @@ int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 err_free_vmaster:
 	kfree(state->vmaster);
 err_unprepare_invs:
+	if (!refcount_read(&state->new_domain_invst.tag.users))
+		arm_smmu_iotlb_tag_free(&state->new_domain_invst.tag);
 	kfree(state->new_domain_invst.new_invs);
 	return ret;
 }
-- 
2.43.0

On Wed, Jan 21, 2026 at 05:24:22PM -0800, Nicolin Chen wrote:
> @@ -3220,6 +3241,9 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
>  		break;
>  	}
>  
> +	/* Set a default users counter */
> +	refcount_set(&cur->users, 1);

I think abusing users here is a little too hard to read..

Can we just keep track in the state somehow with a flag?

Or maybe union in a "bool needs_free" that is for the on-stack version
of this structure?

Jason

On Mon, Jan 26, 2026 at 05:06:40PM -0400, Jason Gunthorpe wrote:
> On Wed, Jan 21, 2026 at 05:24:22PM -0800, Nicolin Chen wrote:
> > @@ -3220,6 +3241,9 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
> >  		break;
> >  	}
> >  
> > +	/* Set a default users counter */
> > +	refcount_set(&cur->users, 1);
> 
> I think abusing users here is a little too hard to read..
> 
> Can we just keep track in the state somehow with a flag?
>
> Or maybe union in a "bool needs_free" that is for the on-stack version
> of this structure?

This doesn't only apply to the case when a tag is newly allocated
during the attach, but it can be used when an old tag has users=0
during a detach, right?

Hmm, perhaps I should have just followed your suggestion, letting
arm_smmu_invs_unref() set the users at all:
https://lore.kernel.org/all/20251219170551.GF254720@nvidia.com/

FWIW, I am hoping to squash this into the base series (v10), so we
wouldn't add a free_fn function pointer in the base series and then
remove it in this followup series.

Nicolin

On Mon, Jan 26, 2026 at 02:23:36PM -0800, Nicolin Chen wrote:
> On Mon, Jan 26, 2026 at 05:06:40PM -0400, Jason Gunthorpe wrote:
> > On Wed, Jan 21, 2026 at 05:24:22PM -0800, Nicolin Chen wrote:
> > > @@ -3220,6 +3241,9 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
> > >  		break;
> > >  	}
> > >  
> > > +	/* Set a default users counter */
> > > +	refcount_set(&cur->users, 1);
> > 
> > I think abusing users here is a little too hard to read..
> > 
> > Can we just keep track in the state somehow with a flag?
> >
> > Or maybe union in a "bool needs_free" that is for the on-stack version
> > of this structure?
> 
> This doesn't only apply to the case when a tag is newly allocated
> during the attach, but it can be used when an old tag has users=0
> during a detach, right?

Oh, IDK, I didn't look so closely..

I though this was just about cleaning up the tag during attach, the
detach flow worked differently and did use users = 0?

Jason

On Tue, Jan 27, 2026 at 12:29:08PM -0400, Jason Gunthorpe wrote:
> On Mon, Jan 26, 2026 at 02:23:36PM -0800, Nicolin Chen wrote:
> > On Mon, Jan 26, 2026 at 05:06:40PM -0400, Jason Gunthorpe wrote:
> > > On Wed, Jan 21, 2026 at 05:24:22PM -0800, Nicolin Chen wrote:
> > > > @@ -3220,6 +3241,9 @@ arm_smmu_master_build_inv(struct arm_smmu_master *master,
> > > >  		break;
> > > >  	}
> > > >  
> > > > +	/* Set a default users counter */
> > > > +	refcount_set(&cur->users, 1);
> > > 
> > > I think abusing users here is a little too hard to read..
> > > 
> > > Can we just keep track in the state somehow with a flag?
> > >
> > > Or maybe union in a "bool needs_free" that is for the on-stack version
> > > of this structure?
> > 
> > This doesn't only apply to the case when a tag is newly allocated
> > during the attach, but it can be used when an old tag has users=0
> > during a detach, right?
> 
> Oh, IDK, I didn't look so closely..
> 
> I though this was just about cleaning up the tag during attach, the
> detach flow worked differently and did use users = 0?

Yes. Detach doesn't allocate a new tag, as it gets the old tag
and decrease its users counter in arm_smmu_invs_unref():

+ * This function will not fail. Any entry with users=0 will be marked as trash,
+ * and caller will be notified about the trashed entry via @to_unref by setting
+ * a users=0.

With that being said, we may still add a "bool is_trash", if we
want. I just don't feel very necessary since we check "is_trash"
throughout the driver via the users counter already..

FWIW, I kept three refcount_set calls in the latest base series
v10, copied from your earlier suggestion. I moved that to this
arm_smmu_master_build_inv() as I felt a bit redundant. Yet, it
is probably better for readability reason:

+	arm_smmu_invs_for_each_cmp(invs, i, to_unref, j, cmp) {
+		if (cmp < 0) {
+			/* not found in to_unref, leave alone */
+			WRITE_ONCE(to_unref->inv[j].users, 1);
+			num_invs = i + 1;
+		} else if (cmp == 0) {
+			int users = READ_ONCE(invs->inv[i].users) - 1;
+
+			if (WARN_ON(users < 0))
+				continue;
+
+			/* same item */
+			WRITE_ONCE(invs->inv[i].users, users);
+			if (users) {
+				WRITE_ONCE(to_unref->inv[j].users, 1);
+				num_invs = i + 1;
+				continue;
+			}
+
+			/* Notify the caller about the trash entry */
+			WRITE_ONCE(to_unref->inv[j].users, 0);
+			invs->num_trashes++;
+		} else {
+			/* item in to_unref is not in invs or already a trash */
+			WARN_ON(true);
+		}
+	}

Nicolin