*/
 int iommu_attach_device(struct iommu_domain *domain, struct device *dev)
 {
-       struct iommu_group *group;
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
        int ret;
 
-       group = iommu_group_get(dev);
        if (!group)
                return -ENODEV;
 
 
 out_unlock:
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
-
        return ret;
 }
 EXPORT_SYMBOL_GPL(iommu_attach_device);
 
 void iommu_detach_device(struct iommu_domain *domain, struct device *dev)
 {
-       struct iommu_group *group;
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
 
-       group = iommu_group_get(dev);
        if (!group)
                return;
 
 
 out_unlock:
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
 }
 EXPORT_SYMBOL_GPL(iommu_detach_device);
 
 struct iommu_domain *iommu_get_domain_for_dev(struct device *dev)
 {
-       struct iommu_domain *domain;
-       struct iommu_group *group;
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
 
-       group = iommu_group_get(dev);
        if (!group)
                return NULL;
 
-       domain = group->domain;
-
-       iommu_group_put(group);
-
-       return domain;
+       return group->domain;
 }
 EXPORT_SYMBOL_GPL(iommu_get_domain_for_dev);
 
  */
 int iommu_device_use_default_domain(struct device *dev)
 {
-       struct iommu_group *group = iommu_group_get(dev);
+       /* Caller is the driver core during the pre-probe path */
+       struct iommu_group *group = dev->iommu_group;
        int ret = 0;
 
        if (!group)
 
 unlock_out:
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
-
        return ret;
 }
 
  */
 void iommu_device_unuse_default_domain(struct device *dev)
 {
-       struct iommu_group *group = iommu_group_get(dev);
+       /* Caller is the driver core during the post-probe path */
+       struct iommu_group *group = dev->iommu_group;
 
        if (!group)
                return;
                group->owner_cnt--;
 
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
 }
 
 static int __iommu_group_alloc_blocking_domain(struct iommu_group *group)
  */
 int iommu_device_claim_dma_owner(struct device *dev, void *owner)
 {
-       struct iommu_group *group;
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
        int ret = 0;
 
        if (WARN_ON(!owner))
                return -EINVAL;
 
-       group = iommu_group_get(dev);
        if (!group)
                return -ENODEV;
 
        ret = __iommu_take_dma_ownership(group, owner);
 unlock_out:
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
-
        return ret;
 }
 EXPORT_SYMBOL_GPL(iommu_device_claim_dma_owner);
  */
 void iommu_device_release_dma_owner(struct device *dev)
 {
-       struct iommu_group *group = iommu_group_get(dev);
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
 
        mutex_lock(&group->mutex);
        if (group->owner_cnt > 1)
        else
                __iommu_release_dma_ownership(group);
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
 }
 EXPORT_SYMBOL_GPL(iommu_device_release_dma_owner);
 
 int iommu_attach_device_pasid(struct iommu_domain *domain,
                              struct device *dev, ioasid_t pasid)
 {
-       struct iommu_group *group;
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
        void *curr;
        int ret;
 
        if (!domain->ops->set_dev_pasid)
                return -EOPNOTSUPP;
 
-       group = iommu_group_get(dev);
        if (!group)
                return -ENODEV;
 
        }
 out_unlock:
        mutex_unlock(&group->mutex);
-       iommu_group_put(group);
-
        return ret;
 }
 EXPORT_SYMBOL_GPL(iommu_attach_device_pasid);
 void iommu_detach_device_pasid(struct iommu_domain *domain, struct device *dev,
                               ioasid_t pasid)
 {
-       struct iommu_group *group = iommu_group_get(dev);
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
 
        mutex_lock(&group->mutex);
        __iommu_remove_group_pasid(group, pasid);
        WARN_ON(xa_erase(&group->pasid_array, pasid) != domain);
        mutex_unlock(&group->mutex);
-
-       iommu_group_put(group);
 }
 EXPORT_SYMBOL_GPL(iommu_detach_device_pasid);
 
                                                    ioasid_t pasid,
                                                    unsigned int type)
 {
+       /* Caller must be a probed driver on dev */
+       struct iommu_group *group = dev->iommu_group;
        struct iommu_domain *domain;
-       struct iommu_group *group;
 
-       group = iommu_group_get(dev);
        if (!group)
                return NULL;
 
        if (type && domain && domain->type != type)
                domain = ERR_PTR(-EBUSY);
        xa_unlock(&group->pasid_array);
-       iommu_group_put(group);
 
        return domain;
 }