static int _dpu_kms_mmu_init(struct dpu_kms *dpu_kms)
 {
-       struct iommu_domain *domain;
        struct msm_gem_address_space *aspace;
-       struct msm_mmu *mmu;
-       struct device *dpu_dev = dpu_kms->dev->dev;
-       struct device *mdss_dev = dpu_dev->parent;
-       struct device *iommu_dev;
-
-       domain = iommu_domain_alloc(&platform_bus_type);
-       if (!domain)
-               return 0;
-
-       /*
-        * IOMMUs can be a part of MDSS device tree binding, or the
-        * MDP/DPU device.
-        */
-       if (dev_iommu_fwspec_get(dpu_dev))
-               iommu_dev = dpu_dev;
-       else
-               iommu_dev = mdss_dev;
-
-       mmu = msm_iommu_new(iommu_dev, domain);
-       if (IS_ERR(mmu)) {
-               iommu_domain_free(domain);
-               return PTR_ERR(mmu);
-       }
-       aspace = msm_gem_address_space_create(mmu, "dpu1",
-               0x1000, 0x100000000 - 0x1000);
 
-       if (IS_ERR(aspace)) {
-               mmu->funcs->destroy(mmu);
+       aspace = msm_kms_init_aspace(dpu_kms->dev);
+       if (IS_ERR(aspace))
                return PTR_ERR(aspace);
-       }
 
        dpu_kms->base.aspace = aspace;
+
        return 0;
 }
 
 
        struct msm_kms *kms;
        struct msm_gem_address_space *aspace;
        int irq, i, ret;
-       struct device *iommu_dev;
-       struct iommu_domain *iommu;
 
        ret = mdp5_init(to_platform_device(dev->dev), dev);
 
        }
        mdelay(16);
 
-       iommu = iommu_domain_alloc(&platform_bus_type);
-       if (iommu) {
-               struct msm_mmu *mmu;
-
-               iommu_dev = &pdev->dev;
-               if (!dev_iommu_fwspec_get(iommu_dev))
-                       iommu_dev = iommu_dev->parent;
-
-               mmu = msm_iommu_new(iommu_dev, iommu);
-
-               aspace = msm_gem_address_space_create(mmu, "mdp5",
-                       0x1000, 0x100000000 - 0x1000);
-
-               if (IS_ERR(aspace)) {
-                       if (!IS_ERR(mmu))
-                               mmu->funcs->destroy(mmu);
-                       ret = PTR_ERR(aspace);
-                       goto fail;
-               }
-
-               kms->aspace = aspace;
-       } else {
-               DRM_DEV_INFO(&pdev->dev,
-                        "no iommu, fallback to phys contig buffers for scanout\n");
-               aspace = NULL;
+       aspace = msm_kms_init_aspace(mdp5_kms->dev);
+       if (IS_ERR(aspace)) {
+               ret = PTR_ERR(aspace);
+               goto fail;
        }
 
+       kms->aspace = aspace;
+
        pm_runtime_put_sync(&pdev->dev);
 
        ret = modeset_init(mdp5_kms);
 
 #include "msm_gem.h"
 #include "msm_gpu.h"
 #include "msm_kms.h"
+#include "msm_mmu.h"
 #include "adreno/adreno_gpu.h"
 
 /*
 
 #include <linux/of_address.h>
 
+struct msm_gem_address_space *msm_kms_init_aspace(struct drm_device *dev)
+{
+       struct iommu_domain *domain;
+       struct msm_gem_address_space *aspace;
+       struct msm_mmu *mmu;
+       struct device *mdp_dev = dev->dev;
+       struct device *mdss_dev = mdp_dev->parent;
+       struct device *iommu_dev;
+
+       domain = iommu_domain_alloc(&platform_bus_type);
+       if (!domain) {
+               drm_info(dev, "no IOMMU, fallback to phys contig buffers for scanout\n");
+               return NULL;
+       }
+
+       /*
+        * IOMMUs can be a part of MDSS device tree binding, or the
+        * MDP/DPU device.
+        */
+       if (dev_iommu_fwspec_get(mdp_dev))
+               iommu_dev = mdp_dev;
+       else
+               iommu_dev = mdss_dev;
+
+       mmu = msm_iommu_new(iommu_dev, domain);
+       if (IS_ERR(mmu)) {
+               iommu_domain_free(domain);
+               return ERR_CAST(mmu);
+       }
+
+       aspace = msm_gem_address_space_create(mmu, "mdp_kms",
+               0x1000, 0x100000000 - 0x1000);
+       if (IS_ERR(aspace))
+               mmu->funcs->destroy(mmu);
+
+       return aspace;
+}
+
 bool msm_use_mmu(struct drm_device *dev)
 {
        struct msm_drm_private *priv = dev->dev_private;
 
 int msm_register_mmu(struct drm_device *dev, struct msm_mmu *mmu);
 void msm_unregister_mmu(struct drm_device *dev, struct msm_mmu *mmu);
 
+struct msm_gem_address_space *msm_kms_init_aspace(struct drm_device *dev);
 bool msm_use_mmu(struct drm_device *dev);
 
 int msm_ioctl_gem_submit(struct drm_device *dev, void *data,