diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index 1ed1e14a1f0c..435835058209 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -2723,26 +2723,92 @@ EXPORT_SYMBOL_GPL(iommu_dev_feature_enabled);
  * This should make us safe against a device being attached to a guest as a
  * whole while there are still pasid users on it (aux and sva).
  */
-int iommu_aux_attach_device(struct iommu_domain *domain, struct device *dev)
+int iommu_aux_attach_device(struct iommu_domain *domain,
+			    struct device *phys_dev, struct device *dev)
 {
-	int ret = -ENODEV;
+	struct iommu_group *group;
+	int ret;
 
-	if (domain->ops->aux_attach_dev)
-		ret = domain->ops->aux_attach_dev(domain, dev);
+	if (!domain->ops->aux_attach_dev ||
+	    !iommu_dev_feature_enabled(phys_dev, IOMMU_DEV_FEAT_AUX))
+		return -ENODEV;
 
-	if (!ret)
-		trace_attach_device_to_domain(dev);
+	/* Bare use only. */
+	if (!dev) {
+		ret = domain->ops->aux_attach_dev(domain, phys_dev);
+		if (!ret)
+			trace_attach_device_to_domain(phys_dev);
+
+		return ret;
+	}
+
+	/*
+	 * The caller has created a made-up device (for example, vfio/mdev)
+	 * and allocated an iommu_group for user level direct assignment.
+	 * Make sure that the group has only single device and hasn't been
+	 * attached by any other domain.
+	 */
+	group = iommu_group_get(dev);
+	if (!group)
+		return -ENODEV;
+
+	/*
+	 * Lock the group to make sure the device-count doesn't change while
+	 * we are attaching.
+	 */
+	mutex_lock(&group->mutex);
+	ret = -EINVAL;
+	if ((iommu_group_device_count(group) != 1) || group->domain)
+		goto out_unlock;
+
+	ret = -EBUSY;
+	if (group->default_domain && group->domain != group->default_domain)
+		goto out_unlock;
+
+	ret = domain->ops->aux_attach_dev(domain, phys_dev);
+	if (!ret) {
+		trace_attach_device_to_domain(phys_dev);
+		group->domain = domain;
+	}
+
+out_unlock:
+	mutex_unlock(&group->mutex);
+	iommu_group_put(group);
 
 	return ret;
 }
 EXPORT_SYMBOL_GPL(iommu_aux_attach_device);
 
-void iommu_aux_detach_device(struct iommu_domain *domain, struct device *dev)
+void iommu_aux_detach_device(struct iommu_domain *domain,
+			     struct device *phys_dev, struct device *dev)
 {
-	if (domain->ops->aux_detach_dev) {
-		domain->ops->aux_detach_dev(domain, dev);
-		trace_detach_device_from_domain(dev);
+	struct iommu_group *group;
+
+	if (WARN_ON_ONCE(!domain->ops->aux_detach_dev))
+		return;
+
+	if (!dev) {
+		domain->ops->aux_detach_dev(domain, phys_dev);
+		trace_detach_device_from_domain(phys_dev);
+
+		return;
 	}
+
+	group = iommu_group_get(dev);
+	if (!group)
+		return;
+
+	mutex_lock(&group->mutex);
+	if (WARN_ON(iommu_group_device_count(group) != 1))
+		goto out_unlock;
+
+	domain->ops->aux_detach_dev(domain, phys_dev);
+	group->domain = NULL;
+	trace_detach_device_from_domain(phys_dev);
+
+out_unlock:
+	mutex_unlock(&group->mutex);
+	iommu_group_put(group);
 }
 EXPORT_SYMBOL_GPL(iommu_aux_detach_device);
 
diff --git a/drivers/vfio/vfio_iommu_type1.c b/drivers/vfio/vfio_iommu_type1.c
index 5e556ac9102a..d3be45dfa58e 100644
--- a/drivers/vfio/vfio_iommu_type1.c
+++ b/drivers/vfio/vfio_iommu_type1.c
@@ -1635,7 +1635,8 @@ static int vfio_mdev_attach_domain(struct device *dev, void *data)
 	iommu_device = vfio_mdev_get_iommu_device(dev);
 	if (iommu_device) {
 		if (iommu_dev_feature_enabled(iommu_device, IOMMU_DEV_FEAT_AUX))
-			return iommu_aux_attach_device(domain, iommu_device);
+			return iommu_aux_attach_device(domain,
+						       iommu_device, dev);
 		else
 			return iommu_attach_device(domain, iommu_device);
 	}
@@ -1651,7 +1652,7 @@ static int vfio_mdev_detach_domain(struct device *dev, void *data)
 	iommu_device = vfio_mdev_get_iommu_device(dev);
 	if (iommu_device) {
 		if (iommu_dev_feature_enabled(iommu_device, IOMMU_DEV_FEAT_AUX))
-			iommu_aux_detach_device(domain, iommu_device);
+			iommu_aux_detach_device(domain, iommu_device, dev);
 		else
 			iommu_detach_device(domain, iommu_device);
 	}
diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index 5657d4fef9f2..7da5e67bf7dc 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -632,8 +632,10 @@ bool iommu_dev_has_feature(struct device *dev, enum iommu_dev_features f);
 int iommu_dev_enable_feature(struct device *dev, enum iommu_dev_features f);
 int iommu_dev_disable_feature(struct device *dev, enum iommu_dev_features f);
 bool iommu_dev_feature_enabled(struct device *dev, enum iommu_dev_features f);
-int iommu_aux_attach_device(struct iommu_domain *domain, struct device *dev);
-void iommu_aux_detach_device(struct iommu_domain *domain, struct device *dev);
+int iommu_aux_attach_device(struct iommu_domain *domain,
+			    struct device *phys_dev, struct device *dev);
+void iommu_aux_detach_device(struct iommu_domain *domain,
+			     struct device *phys_dev, struct device *dev);
 int iommu_aux_get_pasid(struct iommu_domain *domain, struct device *dev);
 
 struct iommu_sva *iommu_sva_bind_device(struct device *dev,
@@ -1007,13 +1009,15 @@ iommu_dev_disable_feature(struct device *dev, enum iommu_dev_features feat)
 }
 
 static inline int
-iommu_aux_attach_device(struct iommu_domain *domain, struct device *dev)
+iommu_aux_attach_device(struct iommu_domain *domain,
+			struct device *phys_dev, struct device *dev)
 {
 	return -ENODEV;
 }
 
 static inline void
-iommu_aux_detach_device(struct iommu_domain *domain, struct device *dev)
+iommu_aux_detach_device(struct iommu_domain *domain,
+			struct device *phys_dev, struct device *dev)
 {
 }
 
