#include "etnaviv_gpu.h"
 #include "etnaviv_gem.h"
 #include "etnaviv_mmu.h"
-#include "etnaviv_iommu.h"
-#include "etnaviv_iommu_v2.h"
 #include "common.xml.h"
 #include "state.xml.h"
 #include "state_hi.xml.h"
 int etnaviv_gpu_init(struct etnaviv_gpu *gpu)
 {
        int ret, i;
-       struct iommu_domain *iommu;
-       enum etnaviv_iommu_version version;
-       bool mmuv2;
 
        ret = pm_runtime_get_sync(gpu->dev);
        if (ret < 0) {
                goto fail;
        }
 
-       /* Setup IOMMU.. eventually we will (I think) do this once per context
-        * and have separate page tables per context.  For now, to keep things
-        * simple and to get something working, just use a single address space:
-        */
-       mmuv2 = gpu->identity.minor_features1 & chipMinorFeatures1_MMU_VERSION;
-       dev_dbg(gpu->dev, "mmuv2: %d\n", mmuv2);
-
-       if (!mmuv2) {
-               iommu = etnaviv_iommu_domain_alloc(gpu);
-               version = ETNAVIV_IOMMU_V1;
-       } else {
-               iommu = etnaviv_iommu_v2_domain_alloc(gpu);
-               version = ETNAVIV_IOMMU_V2;
-       }
-
-       if (!iommu) {
-               dev_err(gpu->dev, "Failed to allocate GPU IOMMU domain\n");
-               ret = -ENOMEM;
-               goto fail;
-       }
-
-       gpu->mmu = etnaviv_iommu_new(gpu, iommu, version);
-       if (!gpu->mmu) {
+       gpu->mmu = etnaviv_iommu_new(gpu);
+       if (IS_ERR(gpu->mmu)) {
                dev_err(gpu->dev, "Failed to instantiate GPU IOMMU\n");
-               iommu_domain_free(iommu);
-               ret = -ENOMEM;
+               ret = PTR_ERR(gpu->mmu);
                goto fail;
        }
 
 
        gpu_write(gpu, VIVS_MC_MMU_RA_PAGE_TABLE, pgtable);
 }
 
-struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu)
+struct iommu_domain *etnaviv_iommuv1_domain_alloc(struct etnaviv_gpu *gpu)
 {
        struct etnaviv_iommu_domain *etnaviv_domain;
        int ret;
 
 #include <linux/iommu.h>
 struct etnaviv_gpu;
 
-struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu);
+struct iommu_domain *etnaviv_iommuv1_domain_alloc(struct etnaviv_gpu *gpu);
 void etnaviv_iommuv1_restore(struct etnaviv_gpu *gpu);
-struct iommu_domain *etnaviv_iommu_v2_domain_alloc(struct etnaviv_gpu *gpu);
+struct iommu_domain *etnaviv_iommuv2_domain_alloc(struct etnaviv_gpu *gpu);
 
 #endif /* __ETNAVIV_IOMMU_H__ */
 
 #include "state_hi.xml.h"
 
 
-struct iommu_domain *etnaviv_iommu_v2_domain_alloc(struct etnaviv_gpu *gpu)
+struct iommu_domain *etnaviv_iommuv2_domain_alloc(struct etnaviv_gpu *gpu)
 {
        /* TODO */
        return NULL;
 
  * this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
+#include "common.xml.h"
 #include "etnaviv_drv.h"
 #include "etnaviv_gem.h"
 #include "etnaviv_gpu.h"
        kfree(mmu);
 }
 
-struct etnaviv_iommu *etnaviv_iommu_new(struct etnaviv_gpu *gpu,
-       struct iommu_domain *domain, enum etnaviv_iommu_version version)
+struct etnaviv_iommu *etnaviv_iommu_new(struct etnaviv_gpu *gpu)
 {
+       enum etnaviv_iommu_version version;
        struct etnaviv_iommu *mmu;
 
        mmu = kzalloc(sizeof(*mmu), GFP_KERNEL);
        if (!mmu)
                return ERR_PTR(-ENOMEM);
 
-       mmu->domain = domain;
+       if (!(gpu->identity.minor_features1 & chipMinorFeatures1_MMU_VERSION)) {
+               mmu->domain = etnaviv_iommuv1_domain_alloc(gpu);
+               version = ETNAVIV_IOMMU_V1;
+       } else {
+               mmu->domain = etnaviv_iommuv2_domain_alloc(gpu);
+               version = ETNAVIV_IOMMU_V2;
+       }
+
+       if (!mmu->domain) {
+               dev_err(gpu->dev, "Failed to allocate GPU IOMMU domain\n");
+               kfree(mmu);
+               return ERR_PTR(-ENOMEM);
+       }
+
        mmu->gpu = gpu;
        mmu->version = version;
        mutex_init(&mmu->lock);
        INIT_LIST_HEAD(&mmu->mappings);
 
-       drm_mm_init(&mmu->mm, domain->geometry.aperture_start,
-                   domain->geometry.aperture_end -
-                     domain->geometry.aperture_start + 1);
+       drm_mm_init(&mmu->mm, mmu->domain->geometry.aperture_start,
+                   mmu->domain->geometry.aperture_end -
+                   mmu->domain->geometry.aperture_start + 1);
 
-       iommu_set_fault_handler(domain, etnaviv_fault_handler, gpu->dev);
+       iommu_set_fault_handler(mmu->domain, etnaviv_fault_handler, gpu->dev);
 
        return mmu;
 }
 
 size_t etnaviv_iommu_dump_size(struct etnaviv_iommu *iommu);
 void etnaviv_iommu_dump(struct etnaviv_iommu *iommu, void *buf);
 
-struct etnaviv_iommu *etnaviv_iommu_new(struct etnaviv_gpu *gpu,
-       struct iommu_domain *domain, enum etnaviv_iommu_version version);
+struct etnaviv_iommu *etnaviv_iommu_new(struct etnaviv_gpu *gpu);
 void etnaviv_iommu_restore(struct etnaviv_gpu *gpu);
 
 #endif /* __ETNAVIV_MMU_H__ */