diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index b1dcb1b9b170..90be3efd0e91 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -406,6 +406,7 @@ struct iommu_fault_param {
  * @iopf_param:	 I/O Page Fault queue and data
  * @fwspec:	 IOMMU fwspec data
  * @iommu_dev:	 IOMMU device this device is linked to
+ * @pasid_array: pasid-indexed array of domains attached to pasid
  * @priv:	 IOMMU Driver private data
  * @max_pasids:  number of PASIDs this device can consume
  * @attach_deferred: the dma domain attachment is deferred
@@ -420,6 +421,7 @@ struct dev_iommu {
 	struct iopf_device_param	*iopf_param;
 	struct iommu_fwspec		*fwspec;
 	struct iommu_device		*iommu_dev;
+	struct xarray			pasid_array;
 	void				*priv;
 	u32				max_pasids;
 	u32				attach_deferred:1;
diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index 1a8fb30341e6..7fd53b9e9754 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -48,7 +48,6 @@ struct iommu_group {
 	struct kobject kobj;
 	struct kobject *devices_kobj;
 	struct list_head devices;
-	struct xarray pasid_array;
 	struct mutex mutex;
 	void *iommu_data;
 	void (*iommu_data_release)(void *iommu_data);
@@ -302,6 +301,7 @@ static struct dev_iommu *dev_iommu_get(struct device *dev)
 		return NULL;
 
 	mutex_init(&param->lock);
+	xa_init(&param->pasid_array);
 	dev->iommu = param;
 	return param;
 }
@@ -891,7 +891,6 @@ struct iommu_group *iommu_group_alloc(void)
 	mutex_init(&group->mutex);
 	INIT_LIST_HEAD(&group->devices);
 	INIT_LIST_HEAD(&group->entry);
-	xa_init(&group->pasid_array);
 
 	ret = ida_alloc(&iommu_group_ida, GFP_KERNEL);
 	if (ret < 0) {
@@ -3041,8 +3040,11 @@ static bool iommu_is_default_domain(struct iommu_group *group)
  */
 static void assert_pasid_dma_ownership(struct iommu_group *group)
 {
+	struct group_device *device;
+
 	lockdep_assert_held(&group->mutex);
-	WARN_ON(!xa_empty(&group->pasid_array));
+	for_each_group_device(group, device)
+		WARN_ON(!xa_empty(&device->dev->iommu->pasid_array));
 }
 
 /**
@@ -3281,33 +3283,6 @@ bool iommu_group_dma_owner_claimed(struct iommu_group *group)
 }
 EXPORT_SYMBOL_GPL(iommu_group_dma_owner_claimed);
 
-static int __iommu_set_group_pasid(struct iommu_domain *domain,
-				   struct iommu_group *group, ioasid_t pasid)
-{
-	struct group_device *device;
-	int ret = 0;
-
-	for_each_group_device(group, device) {
-		ret = domain->ops->set_dev_pasid(domain, device->dev, pasid);
-		if (ret)
-			break;
-	}
-
-	return ret;
-}
-
-static void __iommu_remove_group_pasid(struct iommu_group *group,
-				       ioasid_t pasid)
-{
-	struct group_device *device;
-	const struct iommu_ops *ops;
-
-	for_each_group_device(group, device) {
-		ops = dev_iommu_ops(device->dev);
-		ops->remove_dev_pasid(device->dev, pasid);
-	}
-}
-
 /*
  * iommu_attach_device_pasid() - Attach a domain to pasid of device
  * @domain: the iommu domain.
@@ -3319,32 +3294,28 @@ static void __iommu_remove_group_pasid(struct iommu_group *group,
 int iommu_attach_device_pasid(struct iommu_domain *domain,
 			      struct device *dev, ioasid_t pasid)
 {
-	struct iommu_group *group;
+	struct dev_iommu *param = dev->iommu;
 	void *curr;
 	int ret;
 
 	if (!domain->ops->set_dev_pasid)
 		return -EOPNOTSUPP;
 
-	group = iommu_group_get(dev);
-	if (!group)
+	if (!param)
 		return -ENODEV;
 
-	mutex_lock(&group->mutex);
-	curr = xa_cmpxchg(&group->pasid_array, pasid, NULL, domain, GFP_KERNEL);
+	mutex_lock(&param->lock);
+	curr = xa_cmpxchg(&param->pasid_array, pasid, NULL, domain, GFP_KERNEL);
 	if (curr) {
 		ret = xa_err(curr) ? : -EBUSY;
 		goto out_unlock;
 	}
 
-	ret = __iommu_set_group_pasid(domain, group, pasid);
-	if (ret) {
-		__iommu_remove_group_pasid(group, pasid);
-		xa_erase(&group->pasid_array, pasid);
-	}
+	ret = domain->ops->set_dev_pasid(domain, dev, pasid);
+	if (ret)
+		xa_erase(&param->pasid_array, pasid);
 out_unlock:
-	mutex_unlock(&group->mutex);
-	iommu_group_put(group);
+	mutex_unlock(&param->lock);
 
 	return ret;
 }
@@ -3362,14 +3333,13 @@ 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);
+	const struct iommu_ops *ops = dev_iommu_ops(dev);
+	struct dev_iommu *param = dev->iommu;
 
-	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);
+	mutex_lock(&param->lock);
+	ops->remove_dev_pasid(dev, pasid);
+	WARN_ON(xa_erase(&param->pasid_array, pasid) != domain);
+	mutex_unlock(&param->lock);
 }
 EXPORT_SYMBOL_GPL(iommu_detach_device_pasid);
 
@@ -3392,18 +3362,12 @@ struct iommu_domain *iommu_get_domain_for_dev_pasid(struct device *dev,
 						    unsigned int type)
 {
 	struct iommu_domain *domain;
-	struct iommu_group *group;
 
-	group = iommu_group_get(dev);
-	if (!group)
-		return NULL;
-
-	xa_lock(&group->pasid_array);
-	domain = xa_load(&group->pasid_array, pasid);
+	xa_lock(&dev->iommu->pasid_array);
+	domain = xa_load(&dev->iommu->pasid_array, pasid);
 	if (type && domain && domain->type != type)
 		domain = ERR_PTR(-EBUSY);
-	xa_unlock(&group->pasid_array);
-	iommu_group_put(group);
+	xa_unlock(&dev->iommu->pasid_array);
 
 	return domain;
 }
