diff --git a/drivers/staging/etnaviv/etnaviv_gpu.c b/drivers/staging/etnaviv/etnaviv_gpu.c
index dc02c69512ff..c230da7dd526 100644
--- a/drivers/staging/etnaviv/etnaviv_gpu.c
+++ b/drivers/staging/etnaviv/etnaviv_gpu.c
@@ -451,7 +451,8 @@ static void etnaviv_gpu_hw_init(struct etnaviv_gpu *gpu)
 	gpu_write(gpu, VIVS_MC_MEMORY_BASE_ADDR_PEZ, gpu->memory_base);
 	gpu_write(gpu, VIVS_MC_MEMORY_BASE_ADDR_PE, gpu->memory_base);
 
-	/* FIXME: we need to program the GPU table pointer(s) here */
+	/* setup the MMU page table pointers */
+	etnaviv_iommu_domain_restore(gpu, gpu->mmu->domain);
 
 	/* Start command processor */
 	words = etnaviv_buffer_init(gpu);
@@ -479,6 +480,13 @@ int etnaviv_gpu_init(struct etnaviv_gpu *gpu)
 		return ret;
 
 	etnaviv_hw_identify(gpu);
+
+	if (gpu->identity.model == 0) {
+		dev_err(gpu->dev, "Unknown GPU model\n");
+		pm_runtime_put_autosuspend(gpu->dev);
+		return -ENXIO;
+	}
+
 	ret = etnaviv_hw_reset(gpu);
 	if (ret)
 		goto fail;
diff --git a/drivers/staging/etnaviv/etnaviv_iommu.c b/drivers/staging/etnaviv/etnaviv_iommu.c
index 71f94dac650b..5735319215a3 100644
--- a/drivers/staging/etnaviv/etnaviv_iommu.c
+++ b/drivers/staging/etnaviv/etnaviv_iommu.c
@@ -1,6 +1,6 @@
 /*
  * Copyright (C) 2014 Christian Gmeiner <christian.gmeiner@gmail.com>
-  *
+ *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms of the GNU General Public License version 2 as published by
  * the Free Software Foundation.
@@ -20,11 +20,16 @@
 #include <linux/slab.h>
 #include <linux/dma-mapping.h>
 #include <linux/bitops.h>
+#include <linux/version.h>
 
 #include "etnaviv_gpu.h"
 #include "etnaviv_iommu.h"
 #include "state_hi.xml.h"
 
+#if LINUX_VERSION_CODE < KERNEL_VERSION(4,1,0)
+#define OLD_IOMMU
+#endif
+
 #define PT_SIZE		SZ_512K
 #define PT_ENTRIES	(PT_SIZE / sizeof(uint32_t))
 
@@ -36,12 +41,19 @@ struct etnaviv_iommu_domain_pgtable {
 };
 
 struct etnaviv_iommu_domain {
+	struct iommu_domain domain;
+	struct device *dev;
 	void *bad_page_cpu;
 	dma_addr_t bad_page_dma;
 	struct etnaviv_iommu_domain_pgtable pgtable;
 	spinlock_t map_lock;
 };
 
+static struct etnaviv_iommu_domain *to_etnaviv_domain(struct iommu_domain *domain)
+{
+	return container_of(domain, struct etnaviv_iommu_domain, domain);
+}
+
 static int pgtable_alloc(struct etnaviv_iommu_domain_pgtable *pgtable,
 			 size_t size)
 {
@@ -79,62 +91,72 @@ static void pgtable_write(struct etnaviv_iommu_domain_pgtable *pgtable,
 	pgtable->pgtable[index] = paddr;
 }
 
-static int etnaviv_iommu_domain_init(struct iommu_domain *domain)
+static int __etnaviv_iommu_init(struct etnaviv_iommu_domain *etnaviv_domain)
 {
-	struct etnaviv_iommu_domain *etnaviv_domain;
 	uint32_t iova, *p;
 	int ret, i;
 
-	etnaviv_domain = kmalloc(sizeof(*etnaviv_domain), GFP_KERNEL);
-	if (!etnaviv_domain)
-		return -ENOMEM;
-
-	etnaviv_domain->bad_page_cpu = dma_alloc_coherent(NULL, SZ_4K,
+	etnaviv_domain->bad_page_cpu = dma_alloc_coherent(etnaviv_domain->dev,
+						  SZ_4K,
 						  &etnaviv_domain->bad_page_dma,
 						  GFP_KERNEL);
-	if (!etnaviv_domain->bad_page_cpu) {
-		kfree(etnaviv_domain);
+	if (!etnaviv_domain->bad_page_cpu)
 		return -ENOMEM;
-	}
+
 	p = etnaviv_domain->bad_page_cpu;
 	for (i = 0; i < SZ_4K / 4; i++)
 		*p++ = 0xdead55aa;
 
 	ret = pgtable_alloc(&etnaviv_domain->pgtable, PT_SIZE);
 	if (ret < 0) {
-		dma_free_coherent(NULL, SZ_4K, etnaviv_domain->bad_page_cpu,
+		dma_free_coherent(etnaviv_domain->dev, SZ_4K,
+				  etnaviv_domain->bad_page_cpu,
 				  etnaviv_domain->bad_page_dma);
-		kfree(etnaviv_domain);
 		return ret;
 	}
 
-	for (iova = domain->geometry.aperture_start;
-	     iova < domain->geometry.aperture_end; iova += SZ_4K) {
+	for (iova = etnaviv_domain->domain.geometry.aperture_start;
+	     iova < etnaviv_domain->domain.geometry.aperture_end; iova += SZ_4K) {
 		pgtable_write(&etnaviv_domain->pgtable, iova,
 			      etnaviv_domain->bad_page_dma);
 	}
 
 	spin_lock_init(&etnaviv_domain->map_lock);
-	domain->priv = etnaviv_domain;
+
 	return 0;
 }
 
-static void etnaviv_iommu_domain_destroy(struct iommu_domain *domain)
+static void __etnaviv_iommu_free(struct etnaviv_iommu_domain *etnaviv_domain)
 {
-	struct etnaviv_iommu_domain *etnaviv_domain = domain->priv;
-
 	pgtable_free(&etnaviv_domain->pgtable, PT_SIZE);
 
-	dma_free_coherent(NULL, SZ_4K, etnaviv_domain->bad_page_cpu,
+	dma_free_coherent(etnaviv_domain->dev, SZ_4K,
+			  etnaviv_domain->bad_page_cpu,
 			  etnaviv_domain->bad_page_dma);
+
 	kfree(etnaviv_domain);
+}
+
+#ifdef OLD_IOMMU
+static void etnaviv_iommu_domain_destroy(struct iommu_domain *domain)
+{
+	struct etnaviv_iommu_domain *etnaviv_domain = domain->priv;
+
+	__etnaviv_iommu_free(etnaviv_domain);
+
 	domain->priv = NULL;
 }
+#else
+static void etnaviv_domain_free(struct iommu_domain *domain)
+{
+	__etnaviv_iommu_free(to_etnaviv_domain(domain));
+}
+#endif
 
 static int etnaviv_iommu_map(struct iommu_domain *domain, unsigned long iova,
 	   phys_addr_t paddr, size_t size, int prot)
 {
-	struct etnaviv_iommu_domain *etnaviv_domain = domain->priv;
+	struct etnaviv_iommu_domain *etnaviv_domain = to_etnaviv_domain(domain);
 
 	if (size != SZ_4K)
 		return -EINVAL;
@@ -149,7 +171,7 @@ static int etnaviv_iommu_map(struct iommu_domain *domain, unsigned long iova,
 static size_t etnaviv_iommu_unmap(struct iommu_domain *domain,
 	unsigned long iova, size_t size)
 {
-	struct etnaviv_iommu_domain *etnaviv_domain = domain->priv;
+	struct etnaviv_iommu_domain *etnaviv_domain = to_etnaviv_domain(domain);
 
 	if (size != SZ_4K)
 		return -EINVAL;
@@ -165,41 +187,30 @@ static size_t etnaviv_iommu_unmap(struct iommu_domain *domain,
 static phys_addr_t etnaviv_iommu_iova_to_phys(struct iommu_domain *domain,
 	dma_addr_t iova)
 {
-	struct etnaviv_iommu_domain *etnaviv_domain = domain->priv;
+	struct etnaviv_iommu_domain *etnaviv_domain = to_etnaviv_domain(domain);
 
 	return pgtable_read(&etnaviv_domain->pgtable, iova);
 }
 
 static struct iommu_ops etnaviv_iommu_ops = {
-		.domain_init = etnaviv_iommu_domain_init,
+#ifdef OLD_IOMMU
 		.domain_destroy = etnaviv_iommu_domain_destroy,
+#else
+		.domain_free = etnaviv_domain_free,
+#endif
 		.map = etnaviv_iommu_map,
 		.unmap = etnaviv_iommu_unmap,
 		.iova_to_phys = etnaviv_iommu_iova_to_phys,
 		.pgsize_bitmap = SZ_4K,
 };
 
-struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu)
+void etnaviv_iommu_domain_restore(struct etnaviv_gpu *gpu,
+	struct iommu_domain *domain)
 {
-	struct iommu_domain *domain;
-	struct etnaviv_iommu_domain *etnaviv_domain;
+	struct etnaviv_iommu_domain *etnaviv_domain = to_etnaviv_domain(domain);
 	uint32_t pgtable;
-	int ret;
-
-	domain = kzalloc(sizeof(*domain), GFP_KERNEL);
-	if (!domain)
-		return NULL;
-
-	domain->ops = &etnaviv_iommu_ops;
-	domain->geometry.aperture_start = GPU_MEM_START;
-	domain->geometry.aperture_end = GPU_MEM_START + PT_ENTRIES * SZ_4K;
-
-	ret = domain->ops->domain_init(domain);
-	if (ret)
-		goto out_free;
 
 	/* set page table address in MC */
-	etnaviv_domain = domain->priv;
 	pgtable = (uint32_t)etnaviv_domain->pgtable.paddr;
 
 	gpu_write(gpu, VIVS_MC_MMU_FE_PAGE_TABLE, pgtable);
@@ -207,10 +218,33 @@ struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu)
 	gpu_write(gpu, VIVS_MC_MMU_PE_PAGE_TABLE, pgtable);
 	gpu_write(gpu, VIVS_MC_MMU_PEZ_PAGE_TABLE, pgtable);
 	gpu_write(gpu, VIVS_MC_MMU_RA_PAGE_TABLE, pgtable);
+}
+
+struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu)
+{
+	struct etnaviv_iommu_domain *etnaviv_domain;
+	int ret;
+
+	etnaviv_domain = kzalloc(sizeof(*etnaviv_domain), GFP_KERNEL);
+	if (!etnaviv_domain)
+		return NULL;
 
-	return domain;
+	etnaviv_domain->dev = gpu->dev;
+
+#ifndef OLD_IOMMU
+	etnaviv_domain->domain.type = __IOMMU_DOMAIN_PAGING;
+#endif
+	etnaviv_domain->domain.ops = &etnaviv_iommu_ops;
+	etnaviv_domain->domain.geometry.aperture_start = GPU_MEM_START;
+	etnaviv_domain->domain.geometry.aperture_end = GPU_MEM_START + PT_ENTRIES * SZ_4K;
+
+	ret = __etnaviv_iommu_init(etnaviv_domain);
+	if (ret)
+		goto out_free;
+
+	return &etnaviv_domain->domain;
 
 out_free:
-	kfree(domain);
+	kfree(etnaviv_domain);
 	return NULL;
 }
diff --git a/drivers/staging/etnaviv/etnaviv_iommu.h b/drivers/staging/etnaviv/etnaviv_iommu.h
index c0c359d4f166..cf45503f6b6f 100644
--- a/drivers/staging/etnaviv/etnaviv_iommu.h
+++ b/drivers/staging/etnaviv/etnaviv_iommu.h
@@ -21,6 +21,8 @@
 struct etnaviv_gpu;
 
 struct iommu_domain *etnaviv_iommu_domain_alloc(struct etnaviv_gpu *gpu);
+void etnaviv_iommu_domain_restore(struct etnaviv_gpu *gpu,
+	struct iommu_domain *domain);
 struct iommu_domain *etnaviv_iommu_v2_domain_alloc(struct etnaviv_gpu *gpu);
 
 #endif /* __ETNAVIV_IOMMU_H__ */
