@@ -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;
 }