#include <linux/bitops.h>
 #include <linux/property.h>
 #include <linux/fsl/mc.h>
+#include <linux/module.h>
 #include <trace/events/iommu.h>
 
 static struct kset *iommu_group_kset;
        if (!iommu_get_dev_param(dev))
                return -ENOMEM;
 
+       if (!try_module_get(ops->owner)) {
+               ret = -EINVAL;
+               goto err_free_dev_param;
+       }
+
        ret = ops->add_device(dev);
        if (ret)
-               iommu_free_dev_param(dev);
+               goto err_module_put;
+
+       return 0;
 
+err_module_put:
+       module_put(ops->owner);
+err_free_dev_param:
+       iommu_free_dev_param(dev);
        return ret;
 }
 
        if (dev->iommu_group)
                ops->remove_device(dev);
 
-       iommu_free_dev_param(dev);
+       if (dev->iommu_param) {
+               module_put(ops->owner);
+               iommu_free_dev_param(dev);
+       }
 }
 
 static struct iommu_domain *__iommu_domain_alloc(struct bus_type *bus,
 
  * @sva_get_pasid: Get PASID associated to a SVA handle
  * @page_response: handle page request response
  * @cache_invalidate: invalidate translation caches
- * @pgsize_bitmap: bitmap of all possible supported page sizes
  * @sva_bind_gpasid: bind guest pasid and mm
  * @sva_unbind_gpasid: unbind guest pasid and mm
+ * @pgsize_bitmap: bitmap of all possible supported page sizes
+ * @owner: Driver module providing these ops
  */
 struct iommu_ops {
        bool (*capable)(enum iommu_cap);
        int (*sva_unbind_gpasid)(struct device *dev, int pasid);
 
        unsigned long pgsize_bitmap;
+       struct module *owner;
 };
 
 /**