Once arm_smmu_attach_prepare() installs an iotlb tag to smmu_domain->invs,
arm_smmu_domain_get_iotlb_tag() callers should get the tag from the array.

Only the arm_smmu_domain_get_iotlb_tag() caller for new_smmu_domain in the
arm_smmu_attach_prepare_invs() will be allowed to allocate a new tag.

Suggested-by: Jason Gunthorpe <jgg@nvidia.com>
Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h   |  2 +-
 .../iommu/arm/arm-smmu-v3/arm-smmu-v3-sva.c   |  8 ++---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c   | 32 +++++++++++++++++--
 3 files changed, 34 insertions(+), 8 deletions(-)

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 11879148dad0..812314aaaa6a 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1137,7 +1137,7 @@ struct arm_smmu_attach_state {
 
 int arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
 				  struct arm_smmu_device *smmu,
-				  struct arm_smmu_inv *tag);
+				  struct arm_smmu_inv *tag, bool alloc);
 
 int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 			    struct iommu_domain *new_domain);
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-sva.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-sva.c
index dff494584008..5c8960d31a9b 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-sva.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-sva.c
@@ -30,8 +30,8 @@ arm_smmu_update_s1_domain_cd_entry(struct arm_smmu_domain *smmu_domain)
 		if (WARN_ON(!cdptr))
 			continue;
 
-		if (WARN_ON(arm_smmu_domain_get_iotlb_tag(smmu_domain,
-							  master->smmu, &tag)))
+		if (WARN_ON(arm_smmu_domain_get_iotlb_tag(
+			    smmu_domain, master->smmu, &tag, false)))
 			continue;
 		if (WARN_ON(tag.type != INV_TYPE_S1_ASID))
 			continue;
@@ -170,8 +170,8 @@ static void arm_smmu_mm_release(struct mmu_notifier *mn, struct mm_struct *mm)
 		if (WARN_ON(!cdptr))
 			continue;
 
-		if (WARN_ON(arm_smmu_domain_get_iotlb_tag(smmu_domain,
-							  master->smmu, &tag)))
+		if (WARN_ON(arm_smmu_domain_get_iotlb_tag(
+			    smmu_domain, master->smmu, &tag, false)))
 			continue;
 		if (WARN_ON(tag.type != INV_TYPE_S1_ASID))
 			continue;
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 cf0543f276f3..1927eb794db9 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -3120,10 +3120,31 @@ static void arm_smmu_disable_iopf(struct arm_smmu_master *master,
 		iopf_queue_remove_device(master->smmu->evtq.iopf, master->dev);
 }
 
+static int __arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
+					   struct arm_smmu_inv *tag)
+{
+	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 == tag->type &&
+		    invs->inv[i].smmu == tag->smmu &&
+		    refcount_read(&invs->inv[i].users)) {
+			*tag = invs->inv[i];
+			return 0;
+		}
+	}
+
+	return -ENOENT;
+}
+
 int arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
 				  struct arm_smmu_device *smmu,
-				  struct arm_smmu_inv *tag)
+				  struct arm_smmu_inv *tag, bool alloc)
 {
+	int ret;
+
 	/* Decide the type of the iotlb cache tag */
 	switch (smmu_domain->stage) {
 	case ARM_SMMU_DOMAIN_SVA:
@@ -3139,6 +3160,11 @@ int arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
 
 	tag->smmu = smmu;
 
+	/* Re-use an existing IOTLB cache tag in invs (users counter != 0) */
+	ret = __arm_smmu_domain_get_iotlb_tag(smmu_domain, tag);
+	if (!ret || !alloc)
+		return ret;
+
 	if (tag->type == INV_TYPE_S1_ASID)
 		tag->id = smmu_domain->cd.asid;
 	else
@@ -3314,7 +3340,7 @@ static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
 			lockdep_is_held(&arm_smmu_asid_lock));
 
 		ret = arm_smmu_domain_get_iotlb_tag(new_smmu_domain, smmu,
-						    &invst->tag);
+						    &invst->tag, true);
 		if (ret)
 			return ret;
 
@@ -3344,7 +3370,7 @@ static int arm_smmu_attach_prepare_invs(struct arm_smmu_attach_state *state,
 				lockdep_is_held(&arm_smmu_asid_lock));
 
 		ret = arm_smmu_domain_get_iotlb_tag(old_smmu_domain, smmu,
-						    &invst->tag);
+						    &invst->tag, false);
 		if (WARN_ON(ret))
 			return ret;
 
-- 
2.43.0

On Wed, Jan 21, 2026 at 05:24:21PM -0800, Nicolin Chen wrote:
> Once arm_smmu_attach_prepare() installs an iotlb tag to smmu_domain->invs,
> arm_smmu_domain_get_iotlb_tag() callers should get the tag from the array.
> 
> Only the arm_smmu_domain_get_iotlb_tag() caller for new_smmu_domain in the
> arm_smmu_attach_prepare_invs() will be allowed to allocate a new tag.

It would be nicer to have two functions 'alloc' and 'find' instead of
the bool. Alloc would call find.

The logic looks sound otherwise

Jason

On Mon, Jan 26, 2026 at 05:03:01PM -0400, Jason Gunthorpe wrote:
> On Wed, Jan 21, 2026 at 05:24:21PM -0800, Nicolin Chen wrote:
> > Once arm_smmu_attach_prepare() installs an iotlb tag to smmu_domain->invs,
> > arm_smmu_domain_get_iotlb_tag() callers should get the tag from the array.
> > 
> > Only the arm_smmu_domain_get_iotlb_tag() caller for new_smmu_domain in the
> > arm_smmu_attach_prepare_invs() will be allowed to allocate a new tag.
> 
> It would be nicer to have two functions 'alloc' and 'find' instead of
> the bool. Alloc would call find.

Yea, that just came to me as well after seeing Will's comments
in another thread.

Nicolin