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 69b1d09fd0284..3c5ff4f58934a 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -1926,7 +1926,7 @@ static void __arm_smmu_tlb_inv_range(struct arm_smmu_cmdq_ent *cmd,
 
 	if (smmu->features & ARM_SMMU_FEAT_RANGE_INV) {
 		/* Get the leaf page size */
-		tg = __ffs(smmu_domain->domain.pgsize_bitmap);
+		tg = __ffs(smmu_domain->smmu->pgsize_bitmap);
 
 		/* Convert page size of 12,14,16 (log2) to 1,2,3 */
 		cmd->tlbi.tg = (tg - 10) / 2;
@@ -2053,6 +2053,14 @@ static bool arm_smmu_capable(struct device *dev, enum iommu_cap cap)
 	}
 }
 
+static void arm_smmu_init_smmu_domain(struct arm_smmu_domain *smmu_domain)
+{
+	mutex_init(&smmu_domain->init_mutex);
+	INIT_LIST_HEAD(&smmu_domain->attached_domains);
+	spin_lock_init(&smmu_domain->attached_domains_lock);
+	INIT_LIST_HEAD(&smmu_domain->mmu_notifiers);
+}
+
 static struct iommu_domain *arm_smmu_domain_alloc(unsigned type)
 {
 	struct arm_smmu_domain *smmu_domain;
@@ -2075,14 +2083,22 @@ static struct iommu_domain *arm_smmu_domain_alloc(unsigned type)
 	if (!smmu_domain)
 		return NULL;
 
-	mutex_init(&smmu_domain->init_mutex);
-	INIT_LIST_HEAD(&smmu_domain->attached_domains);
-	spin_lock_init(&smmu_domain->attached_domains_lock);
-	INIT_LIST_HEAD(&smmu_domain->mmu_notifiers);
-
+	arm_smmu_init_smmu_domain(smmu_domain);
 	return &smmu_domain->domain;
 }
 
+struct arm_smmu_domain *
+arm_smmu_init_shared_cd_domain(struct arm_smmu_device *smmu,
+			       struct arm_smmu_domain *smmu_domain,
+			       struct arm_smmu_ctx_desc *cd)
+{
+	arm_smmu_init_smmu_domain(smmu_domain);
+	smmu_domain->smmu = smmu;
+	smmu_domain->stage = ARM_SMMU_DOMAIN_S1_SHARED_CD;
+	smmu_domain->shared_cd = cd;
+	return smmu_domain;
+}
+
 static int arm_smmu_bitmap_alloc(unsigned long *map, int span)
 {
 	int idx, size = 1 << span;
@@ -2541,11 +2557,9 @@ static int arm_smmu_set_dev_pasid(struct iommu_domain *domain,
 				  struct device *dev, ioasid_t pasid)
 {
 	int ret = 0;
-	unsigned long flags;
 	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
 	struct arm_smmu_device *smmu;
 	struct arm_smmu_domain *smmu_domain = to_smmu_domain(domain);
-	struct arm_smmu_attached_domain *attached_domain;
 	struct arm_smmu_master *master;
 
 	if (!fwspec)
@@ -2558,6 +2572,18 @@ static int arm_smmu_set_dev_pasid(struct iommu_domain *domain,
 	if (ret)
 		return ret;
 
+	return arm_smmu_domain_set_dev_pasid(dev, master, smmu_domain, pasid);
+}
+
+int arm_smmu_domain_set_dev_pasid(struct device *dev,
+				  struct arm_smmu_master *master,
+				  struct arm_smmu_domain *smmu_domain,
+				  ioasid_t pasid)
+{
+	unsigned long flags;
+	struct arm_smmu_attached_domain *attached_domain;
+	int ret;
+
 	if (pasid == 0) {
 		dev_err(dev, "pasid 0 is reserved for the device's primary domain\n");
 		return -ENODEV;
@@ -2991,12 +3017,8 @@ static int arm_smmu_def_domain_type(struct device *dev)
 static void arm_smmu_remove_dev_pasid(struct device *dev, ioasid_t pasid)
 {
 	struct iommu_domain *domain;
-	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
-	struct arm_smmu_domain *smmu_domain;
-	struct arm_smmu_attached_domain *attached_domain;
-	unsigned long flags;
 
-	if (!master || pasid == 0)
+	if (pasid == 0)
 		return;
 
 	domain = iommu_get_domain_for_dev_pasid(dev, pasid, 0);
@@ -3005,7 +3027,20 @@ static void arm_smmu_remove_dev_pasid(struct device *dev, ioasid_t pasid)
 	if (domain->type == IOMMU_DOMAIN_SVA)
 		return arm_smmu_sva_remove_dev_pasid(domain, dev, pasid);
 
-	smmu_domain = to_smmu_domain(domain);
+	arm_smmu_domain_remove_dev_pasid(dev, to_smmu_domain(domain), pasid);
+}
+
+void arm_smmu_domain_remove_dev_pasid(struct device *dev,
+				      struct arm_smmu_domain *smmu_domain,
+				      ioasid_t pasid)
+{
+	struct arm_smmu_attached_domain *attached_domain;
+	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
+	unsigned long flags;
+
+	if (!master)
+		return;
+
 	mutex_lock(&arm_smmu_asid_lock);
 	spin_lock_irqsave(&smmu_domain->attached_domains_lock, flags);
 	list_for_each_entry(attached_domain, &smmu_domain->attached_domains, domain_head) {
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 4ac69427abf1c..2c33c0461036d 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -761,6 +761,17 @@ void arm_smmu_tlb_inv_range_asid(unsigned long iova, size_t size, int asid,
 bool arm_smmu_free_asid(struct arm_smmu_ctx_desc *cd);
 int arm_smmu_atc_inv_domain_ssid(struct arm_smmu_domain *smmu_domain, int ssid,
 				 unsigned long iova, size_t size);
+int arm_smmu_domain_set_dev_pasid(struct device *dev,
+				  struct arm_smmu_master *master,
+				  struct arm_smmu_domain *smmu_domain,
+				  ioasid_t pasid);
+void arm_smmu_domain_remove_dev_pasid(struct device *dev,
+				      struct arm_smmu_domain *smmu_domain,
+				      ioasid_t pasid);
+struct arm_smmu_domain *
+arm_smmu_init_shared_cd_domain(struct arm_smmu_device *smmu,
+			       struct arm_smmu_domain *smmu_domain,
+			       struct arm_smmu_ctx_desc *cd);
 
 #ifdef CONFIG_ARM_SMMU_V3_SVA
 bool arm_smmu_sva_supported(struct arm_smmu_device *smmu);
