iommu->domain = NULL;
 }
 
-static void sun50i_iommu_detach_device(struct iommu_domain *domain,
-                                      struct device *dev)
+static int sun50i_iommu_identity_attach(struct iommu_domain *identity_domain,
+                                       struct device *dev)
 {
-       struct sun50i_iommu_domain *sun50i_domain = to_sun50i_domain(domain);
        struct sun50i_iommu *iommu = dev_iommu_priv_get(dev);
+       struct sun50i_iommu_domain *sun50i_domain;
 
        dev_dbg(dev, "Detaching from IOMMU domain\n");
 
-       if (iommu->domain != domain)
-               return;
+       if (iommu->domain == identity_domain)
+               return 0;
 
+       sun50i_domain = to_sun50i_domain(iommu->domain);
        if (refcount_dec_and_test(&sun50i_domain->refcnt))
                sun50i_iommu_detach_domain(iommu, sun50i_domain);
+       return 0;
 }
 
+static struct iommu_domain_ops sun50i_iommu_identity_ops = {
+       .attach_dev = sun50i_iommu_identity_attach,
+};
+
+static struct iommu_domain sun50i_iommu_identity_domain = {
+       .type = IOMMU_DOMAIN_IDENTITY,
+       .ops = &sun50i_iommu_identity_ops,
+};
+
 static int sun50i_iommu_attach_device(struct iommu_domain *domain,
                                      struct device *dev)
 {
        if (iommu->domain == domain)
                return 0;
 
-       if (iommu->domain)
-               sun50i_iommu_detach_device(iommu->domain, dev);
+       sun50i_iommu_identity_attach(&sun50i_iommu_identity_domain, dev);
 
        sun50i_iommu_attach_domain(iommu, sun50i_domain);
 
 }
 
 static const struct iommu_ops sun50i_iommu_ops = {
+       .identity_domain = &sun50i_iommu_identity_domain,
        .pgsize_bitmap  = SZ_4K,
        .device_group   = sun50i_iommu_device_group,
        .domain_alloc   = sun50i_iommu_domain_alloc,
        if (!iommu)
                return -ENOMEM;
        spin_lock_init(&iommu->iommu_lock);
+       iommu->domain = &sun50i_iommu_identity_domain;
        platform_set_drvdata(pdev, iommu);
        iommu->dev = &pdev->dev;