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 d130d723cc33..c9b9c7921bee 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -1104,6 +1104,8 @@ int arm_vsmmu_attach_prepare(struct arm_smmu_attach_state *state,
 			     struct arm_vsmmu *vsmmu);
 void arm_smmu_attach_commit_vmaster(struct arm_smmu_attach_state *state);
 void arm_smmu_master_clear_vmaster(struct arm_smmu_master *master);
+void arm_vsmmu_remove_ats_device(struct arm_vsmmu *vsmmu,
+				 struct arm_smmu_master *master);
 int arm_vmaster_report_event(struct arm_smmu_vmaster *vmaster, u64 *evt);
 
 void arm_smmu_s2_parent_tlb_inv_domain(struct arm_smmu_domain *s2_parent);
@@ -1130,6 +1132,11 @@ arm_smmu_master_clear_vmaster(struct arm_smmu_master *master)
 {
 }
 
+static inline void arm_vsmmu_remove_ats_device(struct arm_vsmmu *vsmmu,
+					       struct arm_smmu_master *master)
+{
+}
+
 static inline int arm_vmaster_report_event(struct arm_smmu_vmaster *vmaster,
 					   u64 *evt)
 {
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
index 4730ff56cf04..491f2b88e30b 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3-iommufd.c
@@ -182,11 +182,13 @@ static void arm_smmu_make_nested_domain_ste(
 int arm_vsmmu_attach_prepare(struct arm_smmu_attach_state *state,
 			     struct arm_vsmmu *vsmmu)
 {
+	struct arm_smmu_master *master = state->master;
 	struct arm_smmu_vmaster *vmaster;
+	unsigned long flags;
 	unsigned long vsid;
 	int ret;
 
-	iommu_group_mutex_assert(state->master->dev);
+	iommu_group_mutex_assert(master->dev);
 
 	ret = iommufd_viommu_get_vdev_id(&vsmmu->core,
 					 state->master->dev, &vsid);
@@ -200,6 +202,12 @@ int arm_vsmmu_attach_prepare(struct arm_smmu_attach_state *state,
 	vmaster->vsid = vsid;
 	state->vmaster = vmaster;
 
+	if (state->ats_enabled) {
+		spin_lock_irqsave(&vsmmu->ats_devices.lock, flags);
+		list_add(&master->devices_elm, &vsmmu->ats_devices.list);
+		spin_unlock_irqrestore(&vsmmu->ats_devices.lock, flags);
+	}
+
 	return 0;
 }
 
@@ -220,6 +228,23 @@ void arm_smmu_master_clear_vmaster(struct arm_smmu_master *master)
 	arm_smmu_attach_commit_vmaster(&state);
 }
 
+void arm_vsmmu_remove_ats_device(struct arm_vsmmu *vsmmu,
+				 struct arm_smmu_master *master)
+{
+	struct arm_smmu_cmdq_ent cmd = { .opcode = CMDQ_OP_ATC_INV };
+	struct arm_smmu_cmdq_batch cmds;
+	unsigned long flags;
+
+	arm_smmu_cmdq_batch_init(vsmmu->smmu, &cmds, &cmd);
+
+	spin_lock_irqsave(&vsmmu->ats_devices.lock, flags);
+	list_del(&master->devices_elm);
+	arm_vsmmu_cmdq_batch_add_atc_inv(vsmmu, master, &cmds, &cmd);
+	spin_unlock_irqrestore(&vsmmu->ats_devices.lock, flags);
+
+	arm_smmu_cmdq_batch_submit(vsmmu->smmu, &cmds);
+}
+
 static int arm_smmu_attach_dev_nested(struct iommu_domain *domain,
 				      struct device *dev)
 {
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 df87880e2a29..483ef9e2c6b7 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -2255,6 +2255,10 @@ static void arm_smmu_tlb_inv_context(void *cookie)
 	 * insertion to guarantee those are observed before the TLBI. Do be
 	 * careful, 007.
 	 */
+
+	if (smmu_domain->nest_parent)
+		return arm_smmu_s2_parent_tlb_inv_domain(smmu_domain);
+
 	if (smmu_domain->stage == ARM_SMMU_DOMAIN_S1)
 		arm_smmu_tlb_inv_asid(smmu, smmu_domain->cd.asid);
 	else
@@ -2342,6 +2346,11 @@ static void arm_smmu_tlb_inv_range_domain(unsigned long iova, size_t size,
 		},
 	};
 
+	if (smmu_domain->nest_parent) {
+		return arm_smmu_s2_parent_tlb_inv_range(smmu_domain, iova, size,
+							granule, leaf);
+	}
+
 	if (smmu_domain->stage == ARM_SMMU_DOMAIN_S1) {
 		cmd.opcode	= smmu_domain->smmu->features & ARM_SMMU_FEAT_E2H ?
 				  CMDQ_OP_TLBI_EL2_VA : CMDQ_OP_TLBI_NH_VA;
@@ -2353,15 +2362,6 @@ static void arm_smmu_tlb_inv_range_domain(unsigned long iova, size_t size,
 	__arm_smmu_tlb_inv_range(smmu_domain->smmu, &cmd, iova, size, granule,
 				 &smmu_domain->domain);
 
-	if (smmu_domain->nest_parent) {
-		/*
-		 * When the S2 domain changes all the nested S1 ASIDs have to be
-		 * flushed too.
-		 */
-		cmd.opcode = CMDQ_OP_TLBI_NH_ALL;
-		arm_smmu_cmdq_issue_cmd_with_sync(smmu_domain->smmu, &cmd);
-	}
-
 	/*
 	 * Unfortunately, this can't be leaf-only since we may have
 	 * zapped an entire table.
@@ -2765,8 +2765,11 @@ static void arm_smmu_remove_master_domain(struct arm_smmu_master *master,
 	if (!smmu_domain)
 		return;
 
-	if (domain->type == IOMMU_DOMAIN_NESTED)
-		nested_ats_flush = to_smmu_nested_domain(domain)->enable_ats;
+	if (domain->type == IOMMU_DOMAIN_NESTED &&
+	    to_smmu_nested_domain(domain)->enable_ats) {
+		return arm_vsmmu_remove_ats_device(
+			to_smmu_nested_domain(domain)->vsmmu, master);
+	}
 
 	spin_lock_irqsave(&smmu_domain->devices_lock, flags);
 	master_domain = arm_smmu_find_master_domain(smmu_domain, master, ssid,
@@ -2837,20 +2840,17 @@ int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 				     arm_smmu_ats_supported(master);
 	}
 
-	if (smmu_domain) {
-		if (new_domain->type == IOMMU_DOMAIN_NESTED) {
-			ret = arm_vsmmu_attach_prepare(
-				state,
-				to_smmu_nested_domain(new_domain)->vsmmu);
-			if (ret)
-				return ret;
-		}
+	if (new_domain->type == IOMMU_DOMAIN_NESTED) {
+		struct arm_smmu_nested_domain *nested_domain =
+			to_smmu_nested_domain(new_domain);
 
+		ret = arm_vsmmu_attach_prepare(state, nested_domain->vsmmu);
+		if (ret)
+			return ret;
+	} else if (smmu_domain) {
 		master_domain = kzalloc(sizeof(*master_domain), GFP_KERNEL);
-		if (!master_domain) {
-			kfree(state->vmaster);
+		if (!master_domain)
 			return -ENOMEM;
-		}
 		master_domain->master = master;
 		master_domain->ssid = state->ssid;
 		if (new_domain->type == IOMMU_DOMAIN_NESTED)
@@ -2877,7 +2877,6 @@ int arm_smmu_attach_prepare(struct arm_smmu_attach_state *state,
 			spin_unlock_irqrestore(&smmu_domain->devices_lock,
 					       flags);
 			kfree(master_domain);
-			kfree(state->vmaster);
 			return -EINVAL;
 		}
 
