return pa;
 }
 
-static int mtk_iommu_add_device(struct device *dev)
+static struct iommu_device *mtk_iommu_probe_device(struct device *dev)
 {
        struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
        struct mtk_iommu_data *data;
-       struct iommu_group *group;
 
        if (!fwspec || fwspec->ops != &mtk_iommu_ops)
-               return -ENODEV; /* Not a iommu client device */
+               return ERR_PTR(-ENODEV); /* Not a iommu client device */
 
        data = dev_iommu_priv_get(dev);
-       iommu_device_link(&data->iommu, dev);
 
-       group = iommu_group_get_for_dev(dev);
-       if (IS_ERR(group))
-               return PTR_ERR(group);
-
-       iommu_group_put(group);
-       return 0;
+       return &data->iommu;
 }
 
-static void mtk_iommu_remove_device(struct device *dev)
+static void mtk_iommu_release_device(struct device *dev)
 {
        struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
-       struct mtk_iommu_data *data;
 
        if (!fwspec || fwspec->ops != &mtk_iommu_ops)
                return;
 
-       data = dev_iommu_priv_get(dev);
-       iommu_device_unlink(&data->iommu, dev);
-
-       iommu_group_remove_device(dev);
        iommu_fwspec_free(dev);
 }
 
        .flush_iotlb_all = mtk_iommu_flush_iotlb_all,
        .iotlb_sync     = mtk_iommu_iotlb_sync,
        .iova_to_phys   = mtk_iommu_iova_to_phys,
-       .add_device     = mtk_iommu_add_device,
-       .remove_device  = mtk_iommu_remove_device,
+       .probe_device   = mtk_iommu_probe_device,
+       .release_device = mtk_iommu_release_device,
        .device_group   = mtk_iommu_device_group,
        .of_xlate       = mtk_iommu_of_xlate,
        .pgsize_bitmap  = SZ_4K | SZ_64K | SZ_1M | SZ_16M,