pr_warn("Failed to set up IOMMU for device %s; retaining platform DMA ops\n",
                 dev_name(dev));
 }
+EXPORT_SYMBOL_GPL(iommu_setup_dma_ops);
 
 static struct iommu_dma_msi_page *iommu_dma_get_msi_page(struct device *dev,
                phys_addr_t msi_addr, struct iommu_domain *domain)
 
 #include <linux/amba/bus.h>
 #include <linux/delay.h>
 #include <linux/dma-iommu.h>
+#include <linux/dma-map-ops.h>
 #include <linux/freezer.h>
 #include <linux/interval_tree.h>
 #include <linux/iommu.h>
        return ERR_PTR(ret);
 }
 
+static void viommu_probe_finalize(struct device *dev)
+{
+#ifndef CONFIG_ARCH_HAS_SETUP_DMA_OPS
+       /* First clear the DMA ops in case we're switching from a DMA domain */
+       set_dma_ops(dev, NULL);
+       iommu_setup_dma_ops(dev, 0, U64_MAX);
+#endif
+}
+
 static void viommu_release_device(struct device *dev)
 {
        struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
        .iova_to_phys           = viommu_iova_to_phys,
        .iotlb_sync             = viommu_iotlb_sync,
        .probe_device           = viommu_probe_device,
+       .probe_finalize         = viommu_probe_finalize,
        .release_device         = viommu_release_device,
        .device_group           = viommu_device_group,
        .get_resv_regions       = viommu_get_resv_regions,