* GNU General Public License for more details.
  */
 
-#include <asm/dma-iommu.h>
-
 #include <drm/drmP.h>
 #include <drm/drm_crtc_helper.h>
 #include <drm/drm_fb_helper.h>
 #include <drm/drm_gem_cma_helper.h>
 #include <drm/drm_of.h>
 #include <linux/dma-mapping.h>
+#include <linux/dma-iommu.h>
 #include <linux/pm_runtime.h>
 #include <linux/module.h>
 #include <linux/of_graph.h>
 #include <linux/component.h>
 #include <linux/console.h>
+#include <linux/iommu.h>
 
 #include "rockchip_drm_drv.h"
 #include "rockchip_drm_fb.h"
 int rockchip_drm_dma_attach_device(struct drm_device *drm_dev,
                                   struct device *dev)
 {
-       struct dma_iommu_mapping *mapping = drm_dev->dev->archdata.mapping;
+       struct rockchip_drm_private *private = drm_dev->dev_private;
        int ret;
 
        if (!is_support_iommu)
                return 0;
 
-       ret = dma_set_coherent_mask(dev, DMA_BIT_MASK(32));
-       if (ret)
+       ret = iommu_attach_device(private->domain, dev);
+       if (ret) {
+               dev_err(dev, "Failed to attach iommu device\n");
                return ret;
+       }
 
-       dma_set_max_seg_size(dev, DMA_BIT_MASK(32));
-
-       return arm_iommu_attach_device(dev, mapping);
+       return 0;
 }
 
 void rockchip_drm_dma_detach_device(struct drm_device *drm_dev,
                                    struct device *dev)
 {
+       struct rockchip_drm_private *private = drm_dev->dev_private;
+       struct iommu_domain *domain = private->domain;
+
        if (!is_support_iommu)
                return;
 
-       arm_iommu_detach_device(dev);
+       iommu_detach_device(domain, dev);
 }
 
 int rockchip_register_crtc_funcs(struct drm_crtc *crtc,
                priv->crtc_funcs[pipe]->disable_vblank(crtc);
 }
 
+static int rockchip_drm_init_iommu(struct drm_device *drm_dev)
+{
+       struct rockchip_drm_private *private = drm_dev->dev_private;
+       struct iommu_domain_geometry *geometry;
+       u64 start, end;
+
+       if (!is_support_iommu)
+               return 0;
+
+       private->domain = iommu_domain_alloc(&platform_bus_type);
+       if (!private->domain)
+               return -ENOMEM;
+
+       geometry = &private->domain->geometry;
+       start = geometry->aperture_start;
+       end = geometry->aperture_end;
+
+       DRM_DEBUG("IOMMU context initialized (aperture: %#llx-%#llx)\n",
+                 start, end);
+       drm_mm_init(&private->mm, start, end - start + 1);
+       mutex_init(&private->mm_lock);
+
+       return 0;
+}
+
+static void rockchip_iommu_cleanup(struct drm_device *drm_dev)
+{
+       struct rockchip_drm_private *private = drm_dev->dev_private;
+
+       if (!is_support_iommu)
+               return;
+
+       drm_mm_takedown(&private->mm);
+       iommu_domain_free(private->domain);
+}
+
 static int rockchip_drm_bind(struct device *dev)
 {
        struct drm_device *drm_dev;
        struct rockchip_drm_private *private;
-       struct dma_iommu_mapping *mapping = NULL;
        int ret;
 
        drm_dev = drm_dev_alloc(&rockchip_drm_driver, dev);
 
        rockchip_drm_mode_config_init(drm_dev);
 
-       dev->dma_parms = devm_kzalloc(dev, sizeof(*dev->dma_parms),
-                                     GFP_KERNEL);
-       if (!dev->dma_parms) {
-               ret = -ENOMEM;
+       ret = rockchip_drm_init_iommu(drm_dev);
+       if (ret)
                goto err_config_cleanup;
-       }
-
-       if (is_support_iommu) {
-               /* TODO(djkurtz): fetch the mapping start/size from somewhere */
-               mapping = arm_iommu_create_mapping(&platform_bus_type,
-                                                  0x00000000,
-                                                  SZ_2G);
-               if (IS_ERR(mapping)) {
-                       ret = PTR_ERR(mapping);
-                       goto err_config_cleanup;
-               }
-
-               ret = dma_set_mask_and_coherent(dev, DMA_BIT_MASK(32));
-               if (ret)
-                       goto err_release_mapping;
-
-               dma_set_max_seg_size(dev, DMA_BIT_MASK(32));
-
-               ret = arm_iommu_attach_device(dev, mapping);
-               if (ret)
-                       goto err_release_mapping;
-       }
 
        /* Try to bind all sub drivers. */
        ret = component_bind_all(dev, drm_dev);
        if (ret)
-               goto err_detach_device;
+               goto err_iommu_cleanup;
 
        /* init kms poll for handling hpd */
        drm_kms_helper_poll_init(drm_dev);
        if (ret)
                goto err_fbdev_fini;
 
-       if (is_support_iommu)
-               arm_iommu_release_mapping(mapping);
        return 0;
 err_fbdev_fini:
        rockchip_drm_fbdev_fini(drm_dev);
 err_kms_helper_poll_fini:
        drm_kms_helper_poll_fini(drm_dev);
        component_unbind_all(dev, drm_dev);
-err_detach_device:
-       if (is_support_iommu)
-               arm_iommu_detach_device(dev);
-err_release_mapping:
-       if (is_support_iommu)
-               arm_iommu_release_mapping(mapping);
+err_iommu_cleanup:
+       rockchip_iommu_cleanup(drm_dev);
 err_config_cleanup:
        drm_mode_config_cleanup(drm_dev);
        drm_dev->dev_private = NULL;
        drm_vblank_cleanup(drm_dev);
        drm_kms_helper_poll_fini(drm_dev);
        component_unbind_all(dev, drm_dev);
-       if (is_support_iommu)
-               arm_iommu_detach_device(dev);
+       rockchip_iommu_cleanup(drm_dev);
        drm_mode_config_cleanup(drm_dev);
        drm_dev->dev_private = NULL;
        drm_dev_unregister(drm_dev);