If the IOTLB cache tag is no longer used by any device, flush the entries
tagged with the ASID/VMID, before the tag gets free-ed.

Drop the free_fn callback accordingly since it's the same thing.

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   |  3 +-
 .../iommu/arm/arm-smmu-v3/arm-smmu-v3-test.c  |  6 +--
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c   | 46 +++++--------------
 3 files changed, 16 insertions(+), 39 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 812314aaaa6a..696ebb89ffa3 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1024,8 +1024,7 @@ void arm_smmu_make_sva_cd(struct arm_smmu_cd *target,
 struct arm_smmu_invs *arm_smmu_invs_merge(struct arm_smmu_invs *invs,
 					  struct arm_smmu_invs *to_merge);
 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);
 struct arm_smmu_invs *arm_smmu_invs_purge(struct arm_smmu_invs *invs);
 #endif
 
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-test.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-test.c
index 81551fad727b..d66931b56b46 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-test.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-test.c
@@ -635,7 +635,7 @@ static void arm_smmu_v3_invs_test(struct kunit *test)
 				     results2[0], results2[1]);
 
 	/* Test3: unref invs2 (same array) */
-	arm_smmu_invs_unref(test_a, &invs2, NULL);
+	arm_smmu_invs_unref(test_a, &invs2);
 	arm_smmu_v3_invs_test_verify(test, test_a, ARRAY_SIZE(results3[0]),
 				     results3[0], results3[1]);
 	KUNIT_EXPECT_EQ(test, test_a->num_trashes, 0);
@@ -647,7 +647,7 @@ static void arm_smmu_v3_invs_test(struct kunit *test)
 				     results4[0], results4[1]);
 
 	/* Test5: unref invs1 (same array) */
-	arm_smmu_invs_unref(test_b, &invs1, NULL);
+	arm_smmu_invs_unref(test_b, &invs1);
 	arm_smmu_v3_invs_test_verify(test, test_b, ARRAY_SIZE(results5[0]),
 				     results5[0], results5[1]);
 	KUNIT_EXPECT_EQ(test, test_b->num_trashes, 2);
@@ -659,7 +659,7 @@ static void arm_smmu_v3_invs_test(struct kunit *test)
 				     results6[0], results6[1]);
 
 	/* Test7: unref invs3 (same array) */
-	arm_smmu_invs_unref(test_a, &invs3, NULL);
+	arm_smmu_invs_unref(test_a, &invs3);
 	KUNIT_EXPECT_EQ(test, test_a->num_invs, 0);
 	KUNIT_EXPECT_EQ(test, test_a->num_trashes, 0);
 
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 d10593823353..5a7032081553 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -1168,8 +1168,6 @@ EXPORT_SYMBOL_IF_KUNIT(arm_smmu_invs_merge);
  *                         the user counts without deletions
  * @invs: the base invalidation array
  * @to_unref: an array of invlidations to decrease their user counts
- * @free_fn: A callback function to invoke, when an entry's user count reduces
- *           to 0
  *
  * Return: the number of trash entries in the array, for arm_smmu_invs_purge()
  *
@@ -1188,8 +1186,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;
@@ -1207,9 +1204,6 @@ void arm_smmu_invs_unref(struct arm_smmu_invs *invs,
 				continue;
 			}
 
-			/* 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++;
@@ -3188,6 +3182,16 @@ int arm_smmu_domain_get_iotlb_tag(struct arm_smmu_domain *smmu_domain,
 
 static void arm_smmu_iotlb_tag_free(struct arm_smmu_inv *tag)
 {
+	struct arm_smmu_cmdq_ent cmd = {
+		.opcode = tag->nsize_opcode,
+	};
+
+	if (tag->type == INV_TYPE_S1_ASID)
+		cmd.tlbi.asid = tag->id;
+	else
+		cmd.tlbi.vmid = tag->id;
+	arm_smmu_cmdq_issue_cmd_with_sync(tag->smmu, &cmd);
+
 	if (tag->type == INV_TYPE_S1_ASID)
 		xa_erase(&arm_smmu_asid_xa, tag->id);
 	else if (tag->type == INV_TYPE_S2_VMID)
@@ -3437,31 +3441,6 @@ 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)
@@ -3475,8 +3454,7 @@ 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 (!refcount_read(&invst->new_invs->inv[0].users))
 		arm_smmu_iotlb_tag_free(&invst->tag);
 
-- 
2.43.0

On Wed, Jan 21, 2026 at 05:24:23PM -0800, Nicolin Chen wrote:
>  static void arm_smmu_iotlb_tag_free(struct arm_smmu_inv *tag)
>  {
> +	struct arm_smmu_cmdq_ent cmd = {
> +		.opcode = tag->nsize_opcode,
> +	};
> +
> +	if (tag->type == INV_TYPE_S1_ASID)
> +		cmd.tlbi.asid = tag->id;
> +	else
> +		cmd.tlbi.vmid = tag->id;
> +	arm_smmu_cmdq_issue_cmd_with_sync(tag->smmu, &cmd);

I think in all these places checking the tag->type it is probably a
good idea to not use a catch all else for vmid? We have many tag types
and some should never come to this, or other, functions.

Jason

On Mon, Jan 26, 2026 at 05:08:33PM -0400, Jason Gunthorpe wrote:
> On Wed, Jan 21, 2026 at 05:24:23PM -0800, Nicolin Chen wrote:
> >  static void arm_smmu_iotlb_tag_free(struct arm_smmu_inv *tag)
> >  {
> > +	struct arm_smmu_cmdq_ent cmd = {
> > +		.opcode = tag->nsize_opcode,
> > +	};
> > +
> > +	if (tag->type == INV_TYPE_S1_ASID)
> > +		cmd.tlbi.asid = tag->id;
> > +	else
> > +		cmd.tlbi.vmid = tag->id;
> > +	arm_smmu_cmdq_issue_cmd_with_sync(tag->smmu, &cmd);
> 
> I think in all these places checking the tag->type it is probably a
> good idea to not use a catch all else for vmid? We have many tag types
> and some should never come to this, or other, functions.

Or maybe we can add an assert function?

static inline void arm_smmu_inv_assert_iotlb_tag(struct arm_smmu_inv *inv)
{
	WARN_ON(inv != INV_TYPE_S1_ASID && inv != INV_TYPE_S1_VMID &&
		inv != INV_TYPE_S1_VMID_VSMMU);
}

Nicolin

On Mon, Jan 26, 2026 at 06:56:36PM -0800, Nicolin Chen wrote:
> On Mon, Jan 26, 2026 at 05:08:33PM -0400, Jason Gunthorpe wrote:
> > On Wed, Jan 21, 2026 at 05:24:23PM -0800, Nicolin Chen wrote:
> > >  static void arm_smmu_iotlb_tag_free(struct arm_smmu_inv *tag)
> > >  {
> > > +	struct arm_smmu_cmdq_ent cmd = {
> > > +		.opcode = tag->nsize_opcode,
> > > +	};
> > > +
> > > +	if (tag->type == INV_TYPE_S1_ASID)
> > > +		cmd.tlbi.asid = tag->id;
> > > +	else
> > > +		cmd.tlbi.vmid = tag->id;
> > > +	arm_smmu_cmdq_issue_cmd_with_sync(tag->smmu, &cmd);
> > 
> > I think in all these places checking the tag->type it is probably a
> > good idea to not use a catch all else for vmid? We have many tag types
> > and some should never come to this, or other, functions.
> 
> Or maybe we can add an assert function?
> 
> static inline void arm_smmu_inv_assert_iotlb_tag(struct arm_smmu_inv *inv)
> {
> 	WARN_ON(inv != INV_TYPE_S1_ASID && inv != INV_TYPE_S1_VMID &&
> 		inv != INV_TYPE_S1_VMID_VSMMU);
> }
 
That works too

Jason