diff --git a/drivers/gpu/msm/kgsl.c b/drivers/gpu/msm/kgsl.c
index 2ee6549e767f..454123e4af0f 100644
--- a/drivers/gpu/msm/kgsl.c
+++ b/drivers/gpu/msm/kgsl.c
@@ -4841,6 +4841,11 @@ int kgsl_device_platform_probe(struct kgsl_device *device)
 	if (status)
 		goto error;
 
+	/* This can return -EPROBE_DEFER */
+	status = kgsl_mmu_probe(device);
+	if (status != 0)
+		goto error_pwrctrl_close;
+
 	if (!devm_request_mem_region(&pdev->dev, device->reg_phys,
 				device->reg_len, device->name)) {
 		dev_err(device->dev, "request_mem_region failed\n");
@@ -4868,10 +4873,6 @@ int kgsl_device_platform_probe(struct kgsl_device *device)
 	rwlock_init(&device->context_lock);
 	spin_lock_init(&device->submit_lock);
 
-	status = kgsl_mmu_probe(device);
-	if (status != 0)
-		goto error_pwrctrl_close;
-
 	/* Check to see if our device can perform DMA correctly */
 	status = dma_set_coherent_mask(&pdev->dev, KGSL_DMA_BIT_MASK);
 	if (status)
diff --git a/drivers/gpu/msm/kgsl_gmu.c b/drivers/gpu/msm/kgsl_gmu.c
index d39f2c0fa8ab..acedbfffaff0 100644
--- a/drivers/gpu/msm/kgsl_gmu.c
+++ b/drivers/gpu/msm/kgsl_gmu.c
@@ -288,10 +288,14 @@ static int gmu_iommu_cb_probe(struct gmu_device *gmu,
 		return -ENODEV;
 
 	pdev = of_find_device_by_node(node);
-	of_dma_configure(&pdev->dev, node, true);
-
+	ret = of_dma_configure(&pdev->dev, node, true);
 	of_node_put(node);
 
+	if (ret) {
+		platform_device_put(pdev);
+		return ret;
+	}
+
 	ctx->pdev = pdev;
 	ctx->domain = iommu_domain_alloc(&platform_bus_type);
 	if (ctx->domain == NULL) {
diff --git a/drivers/gpu/msm/kgsl_iommu.c b/drivers/gpu/msm/kgsl_iommu.c
index a2fa15eca997..fec33d6e49b9 100644
--- a/drivers/gpu/msm/kgsl_iommu.c
+++ b/drivers/gpu/msm/kgsl_iommu.c
@@ -2412,6 +2412,7 @@ static int kgsl_iommu_probe_child(struct kgsl_device *device,
 	struct device_node *node = of_find_node_by_name(parent, name);
 	struct platform_device *pdev;
 	struct device_node *phandle;
+	int ret;
 
 	if (!node)
 		return -ENOENT;
@@ -2430,10 +2431,10 @@ static int kgsl_iommu_probe_child(struct kgsl_device *device,
 
 	of_node_put(phandle);
 
-	of_dma_configure(&pdev->dev, node, true);
+	ret = of_dma_configure(&pdev->dev, node, true);
 
 	of_node_put(node);
-	return 0;
+	return ret;
 }
 
 static void iommu_probe_lpac_context(struct kgsl_device *device,
diff --git a/drivers/gpu/msm/kgsl_mmu.c b/drivers/gpu/msm/kgsl_mmu.c
index dd885a68fc91..398275ecb619 100644
--- a/drivers/gpu/msm/kgsl_mmu.c
+++ b/drivers/gpu/msm/kgsl_mmu.c
@@ -646,8 +646,8 @@ int kgsl_mmu_probe(struct kgsl_device *device)
 	 * go for the NOMMU option instead
 	 */
 	ret = kgsl_iommu_probe(device);
-	if (!ret)
-		return 0;
+	if (!ret || ret == -EPROBE_DEFER)
+		return ret;
 
 	/* set up for NOMMU */
 	set_bit(KGSL_MMU_GLOBAL_PAGETABLE, &mmu->features);