platform_set_drvdata(pdev, obj);
 
        if (omap_iommu_can_register(pdev)) {
-               obj->group = iommu_group_alloc();
-               if (IS_ERR(obj->group))
-                       return PTR_ERR(obj->group);
-
                err = iommu_device_sysfs_add(&obj->iommu, obj->dev, NULL,
                                             obj->name);
                if (err)
-                       goto out_group;
+                       return err;
 
                err = iommu_device_register(&obj->iommu, &omap_iommu_ops, &pdev->dev);
                if (err)
                        goto out_sysfs;
+               obj->has_iommu_driver = true;
        }
 
        pm_runtime_enable(obj->dev);
 
 out_sysfs:
        iommu_device_sysfs_remove(&obj->iommu);
-out_group:
-       iommu_group_put(obj->group);
        return err;
 }
 
 {
        struct omap_iommu *obj = platform_get_drvdata(pdev);
 
-       if (obj->group) {
-               iommu_group_put(obj->group);
-               obj->group = NULL;
-
+       if (obj->has_iommu_driver) {
                iommu_device_sysfs_remove(&obj->iommu);
                iommu_device_unregister(&obj->iommu);
        }
 
 }
 
-static struct iommu_group *omap_iommu_device_group(struct device *dev)
-{
-       struct omap_iommu_arch_data *arch_data = dev_iommu_priv_get(dev);
-       struct iommu_group *group = ERR_PTR(-EINVAL);
-
-       if (!arch_data)
-               return ERR_PTR(-ENODEV);
-
-       if (arch_data->iommu_dev)
-               group = iommu_group_ref_get(arch_data->iommu_dev->group);
-
-       return group;
-}
-
 static const struct iommu_ops omap_iommu_ops = {
        .identity_domain = &omap_iommu_identity_domain,
        .domain_alloc_paging = omap_iommu_domain_alloc_paging,
        .probe_device   = omap_iommu_probe_device,
        .release_device = omap_iommu_release_device,
-       .device_group   = omap_iommu_device_group,
+       .device_group   = generic_single_device_group,
        .pgsize_bitmap  = OMAP_IOMMU_PGSIZES,
        .default_domain_ops = &(const struct iommu_domain_ops) {
                .attach_dev     = omap_iommu_attach_dev,