diff --git a/drivers/iommu/arm/arm-smmu/qcom_iommu.c b/drivers/iommu/arm/arm-smmu/qcom_iommu.c
index b30d6c966e2c..1dec28801eac 100644
--- a/drivers/iommu/arm/arm-smmu/qcom_iommu.c
+++ b/drivers/iommu/arm/arm-smmu/qcom_iommu.c
@@ -573,10 +573,11 @@ static int qcom_iommu_of_xlate(struct device *dev, struct of_phandle_args *args)
 		return -EINVAL;
 	}
 
+	/*
+	 * We're simply retrieving the same platform device that called
+	 * iommu_device_register() when it probed, so it must be valid.
+	 */
 	iommu_pdev = of_find_device_by_node(args->np);
-	if (WARN_ON(!iommu_pdev))
-		return -EINVAL;
-
 	qcom_iommu = platform_get_drvdata(iommu_pdev);
 
 	/* make sure the asid specified in dt is valid, so we don't have
diff --git a/drivers/iommu/exynos-iommu.c b/drivers/iommu/exynos-iommu.c
index de324b4eedfe..73df251d5bcb 100644
--- a/drivers/iommu/exynos-iommu.c
+++ b/drivers/iommu/exynos-iommu.c
@@ -630,6 +630,8 @@ static int exynos_sysmmu_probe(struct platform_device *pdev)
 	if (ret)
 		return ret;
 
+	platform_set_drvdata(pdev, data);
+
 	iommu_device_set_ops(&data->iommu, &exynos_iommu_ops);
 	iommu_device_set_fwnode(&data->iommu, &dev->of_node->fwnode);
 
@@ -637,8 +639,6 @@ static int exynos_sysmmu_probe(struct platform_device *pdev)
 	if (ret)
 		return ret;
 
-	platform_set_drvdata(pdev, data);
-
 	__sysmmu_get_version(data);
 	if (PG_ENT_SHIFT < 0) {
 		if (MMU_MAJ_VER(data->version) < 5) {
@@ -1291,14 +1291,11 @@ static int exynos_iommu_of_xlate(struct device *dev,
 	struct exynos_iommu_owner *owner = dev_iommu_priv_get(dev);
 	struct sysmmu_drvdata *data, *entry;
 
-	if (!sysmmu)
-		return -ENODEV;
-
+	/*
+	 * We're simply retrieving the same platform device that called
+	 * iommu_device_register() when it probed, so it must be valid.
+	 */
 	data = platform_get_drvdata(sysmmu);
-	if (!data) {
-		put_device(&sysmmu->dev);
-		return -ENODEV;
-	}
 
 	if (!owner) {
 		owner = kzalloc(sizeof(*owner), GFP_KERNEL);
diff --git a/drivers/iommu/ipmmu-vmsa.c b/drivers/iommu/ipmmu-vmsa.c
index 0f18abda0e20..6be8dea03d97 100644
--- a/drivers/iommu/ipmmu-vmsa.c
+++ b/drivers/iommu/ipmmu-vmsa.c
@@ -725,11 +725,11 @@ static int ipmmu_init_platform_device(struct device *dev,
 				      struct of_phandle_args *args)
 {
 	struct platform_device *ipmmu_pdev;
-
+	/*
+	 * We're simply retrieving the same platform device that called
+	 * iommu_device_register() when it probed, so it must be valid.
+	 */
 	ipmmu_pdev = of_find_device_by_node(args->np);
-	if (!ipmmu_pdev)
-		return -ENODEV;
-
 	dev_iommu_priv_set(dev, platform_get_drvdata(ipmmu_pdev));
 
 	return 0;
@@ -1075,6 +1075,8 @@ static int ipmmu_probe(struct platform_device *pdev)
 		}
 	}
 
+	platform_set_drvdata(pdev, mmu);
+
 	/*
 	 * Register the IPMMU to the IOMMU subsystem in the following cases:
 	 * - R-Car Gen2 IPMMU (all devices registered)
@@ -1105,9 +1107,6 @@ static int ipmmu_probe(struct platform_device *pdev)
 	 * an IOMMU, which only happens when bus_set_iommu() is called in
 	 * ipmmu_init() after the probe function returns.
 	 */
-
-	platform_set_drvdata(pdev, mmu);
-
 	return 0;
 }
 
diff --git a/drivers/iommu/mtk_iommu.c b/drivers/iommu/mtk_iommu.c
index c072cee532c2..46cba18189ec 100644
--- a/drivers/iommu/mtk_iommu.c
+++ b/drivers/iommu/mtk_iommu.c
@@ -520,11 +520,11 @@ static int mtk_iommu_of_xlate(struct device *dev, struct of_phandle_args *args)
 	}
 
 	if (!dev_iommu_priv_get(dev)) {
-		/* Get the m4u device */
+		/*
+		 * We're simply retrieving the same platform device that called
+		 * iommu_device_register() when it probed, so it must be valid.
+		 */
 		m4updev = of_find_device_by_node(args->np);
-		if (WARN_ON(!m4updev))
-			return -EINVAL;
-
 		dev_iommu_priv_set(dev, platform_get_drvdata(m4updev));
 	}
 
diff --git a/drivers/iommu/rockchip-iommu.c b/drivers/iommu/rockchip-iommu.c
index e5d86b7177de..6a047d6963c0 100644
--- a/drivers/iommu/rockchip-iommu.c
+++ b/drivers/iommu/rockchip-iommu.c
@@ -1096,7 +1096,10 @@ static int rk_iommu_of_xlate(struct device *dev,
 	data = devm_kzalloc(dma_dev, sizeof(*data), GFP_KERNEL);
 	if (!data)
 		return -ENOMEM;
-
+	/*
+	 * We're simply retrieving the same platform device that called
+	 * iommu_device_register() when it probed, so it must be valid.
+	 */
 	iommu_dev = of_find_device_by_node(args->np);
 
 	data->iommu = platform_get_drvdata(iommu_dev);
diff --git a/drivers/iommu/sun50i-iommu.c b/drivers/iommu/sun50i-iommu.c
index ea6db1341916..c0b7a5175b83 100644
--- a/drivers/iommu/sun50i-iommu.c
+++ b/drivers/iommu/sun50i-iommu.c
@@ -763,7 +763,10 @@ static int sun50i_iommu_of_xlate(struct device *dev,
 {
 	struct platform_device *iommu_pdev = of_find_device_by_node(args->np);
 	unsigned id = args->args[0];
-
+	/*
+	 * We're simply retrieving the same platform device that called
+	 * iommu_device_register() when it probed, so it must be valid.
+	 */
 	dev_iommu_priv_set(dev, platform_get_drvdata(iommu_pdev));
 
 	return iommu_fwspec_add_ids(dev, &id, 1);
