diff --git a/drivers/gpu/msm/kgsl_gmu.c b/drivers/gpu/msm/kgsl_gmu.c
index f7b28d9a70a3..848f49682276 100644
--- a/drivers/gpu/msm/kgsl_gmu.c
+++ b/drivers/gpu/msm/kgsl_gmu.c
@@ -25,7 +25,7 @@
 
 struct gmu_iommu_context {
 	const char *name;
-	struct device *dev;
+	struct platform_device *pdev;
 	struct iommu_domain *domain;
 };
 
@@ -277,48 +277,57 @@ static struct gmu_memdesc *allocate_gmu_kmem(struct gmu_device *gmu,
 }
 
 static int gmu_iommu_cb_probe(struct gmu_device *gmu,
-		struct gmu_iommu_context *ctx,
-		struct device_node *node)
+		const char *name, struct gmu_iommu_context *ctx,
+		struct device_node *parent, iommu_fault_handler_t handler)
 {
-	struct platform_device *pdev = of_find_device_by_node(node);
-	struct device *dev;
+	struct device_node *node = of_get_child_by_name(parent, name);
+	struct platform_device *pdev;
 	int ret;
 
-	dev = &pdev->dev;
-	of_dma_configure(dev, node, true);
+	if (!node)
+		return -ENODEV;
 
-	ctx->dev = dev;
+	pdev = of_find_device_by_node(node);
+	of_dma_configure(&pdev->dev, node, true);
+
+	of_node_put(node);
+
+	ctx->pdev = pdev;
 	ctx->domain = iommu_domain_alloc(&platform_bus_type);
 	if (ctx->domain == NULL) {
 		dev_err(&gmu->pdev->dev, "gmu iommu fail to alloc %s domain\n",
 			ctx->name);
+		platform_device_put(pdev);
 		return -ENODEV;
 	}
 
-	ret = iommu_attach_device(ctx->domain, dev);
-	if (ret) {
-		dev_err(&gmu->pdev->dev, "gmu iommu fail to attach %s device\n",
-			ctx->name);
-		iommu_domain_free(ctx->domain);
+	ret = iommu_attach_device(ctx->domain, &pdev->dev);
+	if (!ret) {
+		iommu_set_fault_handler(ctx->domain, handler, ctx);
+		return 0;
 	}
 
+	dev_err(&gmu->pdev->dev,
+		"gmu iommu fail to attach %s device\n", ctx->name);
+	iommu_domain_free(ctx->domain);
+	ctx->domain = NULL;
+	platform_device_put(pdev);
+
 	return ret;
 }
 
-static struct {
-	const char *compatible;
-	int index;
-	iommu_fault_handler_t hdlr;
-} cbs[] = {
-	{ "qcom,smmu-gmu-user-cb",
-		GMU_CONTEXT_USER,
-		gmu_user_fault_handler,
-	},
-	{ "qcom,smmu-gmu-kernel-cb",
-		GMU_CONTEXT_KERNEL,
-		gmu_kernel_fault_handler,
-	},
-};
+
+static void gmu_iommu_cb_close(struct gmu_iommu_context *ctx)
+{
+	if (!ctx->domain)
+		return;
+
+	iommu_detach_device(ctx->domain, &ctx->pdev->dev);
+	iommu_domain_free(ctx->domain);
+
+	platform_device_put(ctx->pdev);
+	memset(ctx, 0, sizeof(*ctx));
+}
 
 /*
  * gmu_iommu_init() - probe IOMMU context banks used by GMU
@@ -328,34 +337,17 @@ static struct {
  */
 static int gmu_iommu_init(struct gmu_device *gmu, struct device_node *node)
 {
-	struct device_node *child;
-	struct gmu_iommu_context *ctx = NULL;
-	int ret, i;
+	int ret;
 
-	of_platform_populate(node, NULL, NULL, &gmu->pdev->dev);
+	devm_of_platform_populate(&gmu->pdev->dev);
 
-	for (i = 0; i < ARRAY_SIZE(cbs); i++) {
-		child = of_find_compatible_node(node, NULL, cbs[i].compatible);
-		if (child) {
-			ctx = &gmu_ctx[cbs[i].index];
-			ret = gmu_iommu_cb_probe(gmu, ctx, child);
-			if (ret)
-				return ret;
-			iommu_set_fault_handler(ctx->domain,
-					cbs[i].hdlr, ctx);
-			}
-		}
+	ret = gmu_iommu_cb_probe(gmu, "gmu_user",
+		&gmu_ctx[GMU_CONTEXT_USER], node, gmu_user_fault_handler);
+	if (ret)
+		return ret;
 
-	for (i = 0; i < ARRAY_SIZE(gmu_ctx); i++) {
-		if (gmu_ctx[i].domain == NULL) {
-			dev_err(&gmu->pdev->dev,
-				"Missing GMU %s context bank node\n",
-				gmu_ctx[i].name);
-			return -EINVAL;
-		}
-	}
-
-	return 0;
+	return gmu_iommu_cb_probe(gmu, "gmu_kernel",
+		&gmu_ctx[GMU_CONTEXT_KERNEL], node, gmu_kernel_fault_handler);
 }
 
 /*
@@ -394,14 +386,8 @@ static void gmu_memory_close(struct gmu_device *gmu)
 		clear_bit(i, &gmu->kmem_bitmap);
 	}
 
-	for (i = 0; i < ARRAY_SIZE(gmu_ctx); i++) {
-		ctx = &gmu_ctx[i];
-
-		if (ctx->domain) {
-			iommu_detach_device(ctx->domain, ctx->dev);
-			iommu_domain_free(ctx->domain);
-		}
-	}
+	gmu_iommu_cb_close(&gmu_ctx[GMU_CONTEXT_USER]);
+	gmu_iommu_cb_close(&gmu_ctx[GMU_CONTEXT_KERNEL]);
 }
 
 static enum gmu_mem_type gmu_get_blk_memtype(struct gmu_device *gmu,
@@ -1061,60 +1047,44 @@ static int gmu_reg_probe(struct kgsl_device *device)
 
 	res = platform_get_resource_byname(gmu->pdev, IORESOURCE_MEM,
 			"kgsl_gmu_reg");
-	if (res == NULL) {
-		dev_err(&gmu->pdev->dev,
-			"platform_get_resource kgsl_gmu_reg failed\n");
-		return -EINVAL;
-	}
-
-	if (res->start == 0 || resource_size(res) == 0) {
-		dev_err(&gmu->pdev->dev,
-				"dev %d kgsl_gmu_reg invalid register region\n",
-				gmu->pdev->dev.id);
-		return -EINVAL;
-	}
-
-	gmu->reg_phys = res->start;
-	gmu->reg_len = resource_size(res);
-	device->gmu_core.reg_virt = devm_ioremap(&gmu->pdev->dev,
-			res->start, resource_size(res));
-	if (device->gmu_core.reg_virt == NULL) {
-		dev_err(&gmu->pdev->dev, "kgsl_gmu_reg ioremap failed\n");
+	if (!res) {
+		dev_err(&gmu->pdev->dev, "The GMU register region isn't defined\n");
 		return -ENODEV;
 	}
 
+	device->gmu_core.gmu2gpu_offset = (res->start - device->reg_phys) >> 2;
+	device->gmu_core.reg_len = resource_size(res);
+
+	/*
+	 * We can't use devm_ioremap_resource here because we purposely double
+	 * map the gpu_cc registers for debugging purposes
+	 */
+	device->gmu_core.reg_virt = devm_ioremap(&gmu->pdev->dev, res->start,
+		resource_size(res));
+
+	if (!device->gmu_core.reg_virt) {
+		dev_err(&gmu->pdev->dev, "Unable to map the GMU registers\n");
+		return -ENOMEM;
+	}
+
 	return 0;
 }
 
 static int gmu_regulators_probe(struct gmu_device *gmu,
-		struct device_node *node)
+		struct platform_device *pdev)
 {
-	const char *name;
-	struct property *prop;
-	struct device *dev = &gmu->pdev->dev;
-	int ret = 0;
+	gmu->cx_gdsc = devm_regulator_get(&pdev->dev, "vddcx");
+	if (IS_ERR(gmu->cx_gdsc)) {
+		if (PTR_ERR(gmu->cx_gdsc) != -EPROBE_DEFER)
+			dev_err(&pdev->dev, "Couldn't get the vddcx gdsc\n");
+		return PTR_ERR(gmu->cx_gdsc);
+	}
 
-	of_property_for_each_string(node, "regulator-names", prop, name) {
-		if (!strcmp(name, "vddcx")) {
-			gmu->cx_gdsc = devm_regulator_get(dev, name);
-			if (IS_ERR(gmu->cx_gdsc)) {
-				ret = PTR_ERR(gmu->cx_gdsc);
-				dev_err(dev, "dt: GMU couldn't get CX gdsc\n");
-				gmu->cx_gdsc = NULL;
-				return ret;
-			}
-		} else if (!strcmp(name, "vdd")) {
-			gmu->gx_gdsc = devm_regulator_get(dev, name);
-			if (IS_ERR(gmu->gx_gdsc)) {
-				ret = PTR_ERR(gmu->gx_gdsc);
-				dev_err(dev, "dt: GMU couldn't get GX gdsc\n");
-				gmu->gx_gdsc = NULL;
-				return ret;
-			}
-		} else {
-			dev_err(dev, "dt: Unknown GMU regulator: %s\n", name);
-			return -ENODEV;
-		}
+	gmu->gx_gdsc = devm_regulator_get(&pdev->dev, "vdd");
+	if (IS_ERR(gmu->gx_gdsc)) {
+		if (PTR_ERR(gmu->gx_gdsc) != -EPROBE_DEFER)
+			dev_err(&pdev->dev, "Couldn't get the vdd gdsc\n");
+		return PTR_ERR(gmu->gx_gdsc);
 	}
 
 	return 0;
@@ -1321,7 +1291,7 @@ static int gmu_probe(struct kgsl_device *device, struct platform_device *pdev)
 
 	gmu->pdev = pdev;
 
-	device->gmu_core.ptr = (void *)gmu;
+	device->gmu_core.ptr = gmu;
 	hfi = &gmu->hfi;
 	gmu->load_mode = TCM_BOOT;
 
@@ -1330,11 +1300,11 @@ static int gmu_probe(struct kgsl_device *device, struct platform_device *pdev)
 	set_dma_ops(&gmu->pdev->dev, NULL);
 
 	/* Set up GMU regulators */
-	ret = gmu_regulators_probe(gmu, pdev->dev.of_node);
+	ret = gmu_regulators_probe(gmu, pdev);
 	if (ret)
 		return ret;
 
-	ret = devm_clk_bulk_get_all(&gmu->pdev->dev, &gmu->clks);
+	ret = devm_clk_bulk_get_all(&pdev->dev, &gmu->clks);
 	if (ret < 0)
 		return ret;
 
@@ -1343,8 +1313,9 @@ static int gmu_probe(struct kgsl_device *device, struct platform_device *pdev)
 	/* Get a pointer to the GMU clock */
 	gmu->gmu_clk = kgsl_of_clk_by_name(gmu->clks, gmu->num_clks, "gmu_clk");
 	if (!gmu->gmu_clk) {
-		dev_err(&gmu->pdev->dev, "Couldn't get gmu_clk\n");
-		return -ENODEV;
+		dev_err(&pdev->dev, "Couldn't get gmu_clk\n");
+		ret = -ENODEV;
+		goto error;
 	}
 
 	/* Set up GMU IOMMU and shared memory with GMU */
@@ -1366,10 +1337,6 @@ static int gmu_probe(struct kgsl_device *device, struct platform_device *pdev)
 	if (ret)
 		goto error;
 
-	device->gmu_core.gmu2gpu_offset =
-			(gmu->reg_phys - device->reg_phys) >> 2;
-	device->gmu_core.reg_len = gmu->reg_len;
-
 	/* Initialize HFI and GMU interrupts */
 	hfi->hfi_interrupt_num = kgsl_request_irq(gmu->pdev, "kgsl_hfi_irq",
 		hfi_irq_handler, device);
@@ -1474,9 +1441,6 @@ static int gmu_enable_gdsc(struct gmu_device *gmu)
 {
 	int ret;
 
-	if (IS_ERR_OR_NULL(gmu->cx_gdsc))
-		return 0;
-
 	ret = regulator_enable(gmu->cx_gdsc);
 	if (ret)
 		dev_err(&gmu->pdev->dev,
@@ -1738,14 +1702,8 @@ static void gmu_remove(struct kgsl_device *device)
 {
 	struct adreno_device *adreno_dev = ADRENO_DEVICE(device);
 	struct gmu_device *gmu = KGSL_GMU_DEVICE(device);
-	struct kgsl_hfi *hfi;
 
-	if (gmu == NULL || gmu->pdev == NULL)
-		return;
-
-	hfi = &gmu->hfi;
-
-	tasklet_kill(&hfi->tasklet);
+	tasklet_kill(&gmu->hfi.tasklet);
 
 	gmu_stop(device);
 
diff --git a/drivers/gpu/msm/kgsl_gmu.h b/drivers/gpu/msm/kgsl_gmu.h
index 7c6ca7e03b5c..b6a815d90965 100644
--- a/drivers/gpu/msm/kgsl_gmu.h
+++ b/drivers/gpu/msm/kgsl_gmu.h
@@ -141,8 +141,6 @@ struct icc_path;
 /**
  * struct gmu_device - GMU device structure
  * @ver: GMU Version information
- * @reg_phys: GMU CSR physical address
- * @reg_len: GMU CSR range
  * @gmu_interrupt_num: GMU interrupt number
  * @fw_image: GMU FW image
  * @hfi_mem: pointer to HFI shared memory
@@ -174,8 +172,6 @@ struct gmu_device {
 		u32 hfi;
 	} ver;
 	struct platform_device *pdev;
-	unsigned long reg_phys;
-	unsigned int reg_len;
 	int gmu_interrupt_num;
 	const struct firmware *fw_image;
 	struct gmu_memdesc *hfi_mem;