/* The domain can be NULL only when processing the first attach */
        if (!domain)
                return NULL;
-       if (domain->type & __IOMMU_DOMAIN_PAGING)
+       if ((domain->type & __IOMMU_DOMAIN_PAGING) ||
+           domain->type == IOMMU_DOMAIN_SVA)
                return to_smmu_domain(domain);
        return NULL;
 }
                       struct arm_smmu_domain *smmu_domain, ioasid_t pasid,
                       const struct arm_smmu_cd *cd)
 {
+       struct arm_smmu_attach_state state = {
+               .master = master,
+               /*
+                * For now the core code prevents calling this when a domain is
+                * already attached, no need to set old_domain.
+                */
+               .ssid = pasid,
+       };
        struct arm_smmu_cd *cdptr;
+       int ret;
 
        /* The core code validates pasid */
 
        cdptr = arm_smmu_alloc_cd_ptr(master, pasid);
        if (!cdptr)
                return -ENOMEM;
+
+       mutex_lock(&arm_smmu_asid_lock);
+       ret = arm_smmu_attach_prepare(&state, &smmu_domain->domain);
+       if (ret)
+               goto out_unlock;
+
        arm_smmu_write_cd_entry(master, pasid, cdptr, cd);
-       return 0;
+
+       arm_smmu_attach_commit(&state);
+
+out_unlock:
+       mutex_unlock(&arm_smmu_asid_lock);
+       return ret;
 }
 
 void arm_smmu_remove_pasid(struct arm_smmu_master *master,
                           struct arm_smmu_domain *smmu_domain, ioasid_t pasid)
 {
+       mutex_lock(&arm_smmu_asid_lock);
        arm_smmu_clear_cd(master, pasid);
+       if (master->ats_enabled)
+               arm_smmu_atc_inv_master(master, pasid);
+       arm_smmu_remove_master_domain(master, &smmu_domain->domain, pasid);
+       mutex_unlock(&arm_smmu_asid_lock);
 }
 
 static int arm_smmu_attach_dev_ste(struct iommu_domain *domain,