---

v2:
  - Comment bus driver changes for clarity
  - Use dev->iommu as the now-robust replay condition
  - Drop the device_iommu_mapped() checks in the firmware paths as they
    weren't doing much - we can't replace probe_device_lock just yet...
  
  drivers/acpi/arm64/dma.c        |  5 +++++
  drivers/acpi/scan.c             |  7 -------
  drivers/amba/bus.c              |  3 ++-
  drivers/base/platform.c         |  3 ++-
  drivers/bus/fsl-mc/fsl-mc-bus.c |  3 ++-
  drivers/cdx/cdx.c               |  3 ++-
  drivers/iommu/iommu.c           | 24 +++++++++++++++++++++---
  drivers/iommu/of_iommu.c        |  7 ++++++-
  drivers/of/device.c             |  7 ++++++-
  drivers/pci/pci-driver.c        |  3 ++-
  10 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/drivers/acpi/arm64/dma.c b/drivers/acpi/arm64/dma.c
index 52b2abf88689..f30f138352b7 100644
--- a/drivers/acpi/arm64/dma.c
+++ b/drivers/acpi/arm64/dma.c

@@ -26,6 +26,11 @@ void acpi_arch_dma_setup(struct device *dev)
  	else
  		end = (1ULL << 32) - 1;
  
+	if (dev->dma_range_map) {
+		dev_dbg(dev, "dma_range_map already set\n");
+		return;
+	}
+
  	ret = acpi_dma_get_range(dev, &map);
  	if (!ret && map) {
  		end = dma_range_map_max(map);

diff --git a/drivers/acpi/scan.c b/drivers/acpi/scan.c
index 9f4efa8f75a6..fb1fe9f3b1a3 100644
--- a/drivers/acpi/scan.c
+++ b/drivers/acpi/scan.c

@@ -1632,13 +1632,6 @@ static int acpi_iommu_configure_id(struct device *dev, const u32 *id_in)
  		err = viot_iommu_configure(dev);
  	mutex_unlock(&iommu_probe_device_lock);
  
-	/*
-	 * If we have reason to believe the IOMMU driver missed the initial
-	 * iommu_probe_device() call for dev, replay it to get things in order.
-	 */
-	if (!err && dev->bus)
-		err = iommu_probe_device(dev);
-
  	return err;
  }
  

diff --git a/drivers/amba/bus.c b/drivers/amba/bus.c
index 8ef259b4d037..71482d639a6d 100644
--- a/drivers/amba/bus.c
+++ b/drivers/amba/bus.c

@@ -364,7 +364,8 @@ static int amba_dma_configure(struct device *dev)
  		ret = acpi_dma_configure(dev, attr);
  	}
  
-	if (!ret && !drv->driver_managed_dma) {
+	/* @drv may not be valid when we're called from the IOMMU layer */
+	if (!ret && dev->driver && !drv->driver_managed_dma) {
  		ret = iommu_device_use_default_domain(dev);
  		if (ret)
  			arch_teardown_dma_ops(dev);

diff --git a/drivers/base/platform.c b/drivers/base/platform.c
index 6f2a33722c52..1813cfd0c4bd 100644
--- a/drivers/base/platform.c
+++ b/drivers/base/platform.c

@@ -1451,7 +1451,8 @@ static int platform_dma_configure(struct device *dev)
  		attr = acpi_get_dma_attr(to_acpi_device_node(fwnode));
  		ret = acpi_dma_configure(dev, attr);
  	}
-	if (ret || drv->driver_managed_dma)
+	/* @drv may not be valid when we're called from the IOMMU layer */
+	if (ret || !dev->driver || drv->driver_managed_dma)
  		return ret;
  
  	ret = iommu_device_use_default_domain(dev);

diff --git a/drivers/bus/fsl-mc/fsl-mc-bus.c b/drivers/bus/fsl-mc/fsl-mc-bus.c
index d1f3d327ddd1..a8be8cf246fb 100644
--- a/drivers/bus/fsl-mc/fsl-mc-bus.c
+++ b/drivers/bus/fsl-mc/fsl-mc-bus.c

@@ -153,7 +153,8 @@ static int fsl_mc_dma_configure(struct device *dev)
  	else
  		ret = acpi_dma_configure_id(dev, DEV_DMA_COHERENT, &input_id);
  
-	if (!ret && !mc_drv->driver_managed_dma) {
+	/* @mc_drv may not be valid when we're called from the IOMMU layer */
+	if (!ret && dev->driver && !mc_drv->driver_managed_dma) {
  		ret = iommu_device_use_default_domain(dev);
  		if (ret)
  			arch_teardown_dma_ops(dev);

diff --git a/drivers/cdx/cdx.c b/drivers/cdx/cdx.c
index c573ed2ee71a..780fb0c4adba 100644
--- a/drivers/cdx/cdx.c
+++ b/drivers/cdx/cdx.c

@@ -360,7 +360,8 @@ static int cdx_dma_configure(struct device *dev)
  		return ret;
  	}
  
-	if (!ret && !cdx_drv->driver_managed_dma) {
+	/* @cdx_drv may not be valid when we're called from the IOMMU layer */
+	if (!ret && dev->driver && !cdx_drv->driver_managed_dma) {
  		ret = iommu_device_use_default_domain(dev);
  		if (ret)
  			arch_teardown_dma_ops(dev);

diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index a3b45b84f42b..1cec7074367a 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c

@@ -414,9 +414,21 @@ static int iommu_init_device(struct device *dev)
  	if (!dev_iommu_get(dev))
  		return -ENOMEM;
  	/*
-	 * For FDT-based systems and ACPI IORT/VIOT, drivers register IOMMU
-	 * instances with non-NULL fwnodes, and client devices should have been
-	 * identified with a fwspec by this point. Otherwise, we can currently
+	 * For FDT-based systems and ACPI IORT/VIOT, the common firmware parsing
+	 * is buried in the bus dma_configure path. Properly unpicking that is
+	 * still a big job, so for now just invoke the whole thing. The device
+	 * already having a driver bound means dma_configure has already run and
+	 * either found no IOMMU to wait for, or we're in its replay call right
+	 * now, so either way there's no point calling it again.
+	 */
+	if (!dev->driver && dev->bus->dma_configure) {
+		mutex_unlock(&iommu_probe_device_lock);
+		dev->bus->dma_configure(dev);
+		mutex_lock(&iommu_probe_device_lock);
+	}
+	/*
+	 * At this point, relevant devices either now have a fwspec which will
+	 * match ops registered with a non-NULL fwnode, or we can reasonably
  	 * assume that only one of Intel, AMD, s390, PAMU or legacy SMMUv2 can
  	 * be present, and that any of their registered instances has suitable
  	 * ops for probing, and thus cheekily co-opt the same mechanism.

@@ -426,6 +438,12 @@ static int iommu_init_device(struct device *dev)
  		ret = -ENODEV;
  		goto err_free;
  	}
+	/*
+	 * And if we do now see any replay calls, they would indicate someone
+	 * misusing the dma_configure path outside bus code.
+	 */
+	if (dev->driver)
+		dev_WARN(dev, "late IOMMU probe at driver bind, something fishy here!\n");
  
  	if (!try_module_get(ops->owner)) {
  		ret = -EINVAL;

diff --git a/drivers/iommu/of_iommu.c b/drivers/iommu/of_iommu.c
index e10a68b5ffde..6b989a62def2 100644
--- a/drivers/iommu/of_iommu.c
+++ b/drivers/iommu/of_iommu.c

@@ -155,7 +155,12 @@ int of_iommu_configure(struct device *dev, struct device_node *master_np,
  		dev_iommu_free(dev);
  	mutex_unlock(&iommu_probe_device_lock);
  
-	if (!err && dev->bus)
+	/*
+	 * If we're not on the iommu_probe_device() path (as indicated by the
+	 * initial dev->iommu) then try to simulate it. This should no longer
+	 * happen unless of_dma_configure() is being misused outside bus code.
+	 */
+	if (!err && dev->bus && !dev_iommu_present)
  		err = iommu_probe_device(dev);
  
  	if (err && err != -EPROBE_DEFER)

diff --git a/drivers/of/device.c b/drivers/of/device.c
index edf3be197265..5053e5d532cc 100644
--- a/drivers/of/device.c
+++ b/drivers/of/device.c

@@ -99,6 +99,11 @@ int of_dma_configure_id(struct device *dev, struct device_node *np,
  	bool coherent, set_map = false;
  	int ret;
  
+	if (dev->dma_range_map) {
+		dev_dbg(dev, "dma_range_map already set\n");
+		goto skip_map;
+	}
+
  	if (np == dev->of_node)
  		bus_np = __of_get_dma_parent(np);
  	else

@@ -119,7 +124,7 @@ int of_dma_configure_id(struct device *dev, struct device_node *np,
  		end = dma_range_map_max(map);
  		set_map = true;
  	}
-
+skip_map:
  	/*
  	 * If @dev is expected to be DMA-capable then the bus code that created
  	 * it should have initialised its dma_mask pointer by this point. For

diff --git a/drivers/pci/pci-driver.c b/drivers/pci/pci-driver.c
index f57ea36d125d..4468b7327cab 100644
--- a/drivers/pci/pci-driver.c
+++ b/drivers/pci/pci-driver.c

@@ -1653,7 +1653,8 @@ static int pci_dma_configure(struct device *dev)
  
  	pci_put_host_bridge_device(bridge);
  
-	if (!ret && !driver->driver_managed_dma) {
+	/* @driver may not be valid when we're called from the IOMMU layer */
+	if (!ret && dev->driver && !driver->driver_managed_dma) {
  		ret = iommu_device_use_default_domain(dev);
  		if (ret)
  			arch_teardown_dma_ops(dev);