diff --git a/drivers/vfio/group.c b/drivers/vfio/group.c
index 95b336de8a17..f43a9bd99907 100644
--- a/drivers/vfio/group.c
+++ b/drivers/vfio/group.c
@@ -842,9 +842,6 @@ struct iommu_group *vfio_file_iommu_group(struct file *file)
 	struct vfio_group *group = vfio_group_from_file(file);
 	struct iommu_group *iommu_group = NULL;
 
-	if (!IS_ENABLED(CONFIG_SPAPR_TCE_IOMMU))
-		return NULL;
-
 	if (!group)
 		return NULL;
 
diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 0ec218d9613a..b450c41e2839 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -1634,6 +1634,31 @@ static __always_inline bool kvm_arch_has_assigned_device(struct kvm *kvm)
 }
 #endif
 
+#ifdef __KVM_HAVE_ARCH_ASSIGNED_DEVICE_GROUP
+int kvm_arch_assign_device(struct device *dev);
+int kvm_arch_assign_group(struct iommu_group *group);
+void kvm_arch_reclaim_device(struct device *dev);
+void kvm_arch_reclaim_group(struct iommu_group *group);
+#else
+static inline int kvm_arch_assign_device(struct device *dev)
+{
+	return 0;
+}
+
+static inline int kvm_arch_assign_group(struct iommu_group *group)
+{
+	return 0;
+}
+
+static inline void kvm_arch_reclaim_device(struct device *dev)
+{
+}
+
+static inline void kvm_arch_reclaim_group(struct iommu_group *group)
+{
+}
+#endif
+
 static inline struct rcuwait *kvm_arch_vcpu_get_wait(struct kvm_vcpu *vcpu)
 {
 #ifdef __KVM_HAVE_ARCH_WQP
diff --git a/virt/kvm/vfio.c b/virt/kvm/vfio.c
index 388ae471d258..fd57020a3d8e 100644
--- a/virt/kvm/vfio.c
+++ b/virt/kvm/vfio.c
@@ -80,7 +80,22 @@ static bool kvm_vfio_file_is_valid(struct file *file)
 	return ret;
 }
 
-#ifdef CONFIG_SPAPR_TCE_IOMMU
+static struct device *kvm_vfio_file_get_device(struct file *file)
+{
+	struct device *(*fn)(struct file *file);
+	struct device *dev;
+
+	fn = symbol_get(vfio_file_get_device);
+	if (!fn)
+		return NULL;
+
+	dev = fn(file);
+
+	symbol_put(vfio_file_get_device);
+
+	return dev;
+}
+
 static struct iommu_group *kvm_vfio_file_iommu_group(struct file *file)
 {
 	struct iommu_group *(*fn)(struct file *file);
@@ -97,6 +112,7 @@ static struct iommu_group *kvm_vfio_file_iommu_group(struct file *file)
 	return ret;
 }
 
+#ifdef CONFIG_SPAPR_TCE_IOMMU
 static void kvm_spapr_tce_release_vfio_group(struct kvm *kvm,
 					     struct kvm_vfio_file *kvf)
 {
@@ -140,6 +156,36 @@ static void kvm_vfio_update_coherency(struct kvm_device *dev)
 	}
 }
 
+static int kvm_vfio_assign_file(struct file *file)
+{
+	struct device *dev;
+	struct iommu_group *group;
+
+	dev = kvm_vfio_file_get_device(file);
+	if (dev)
+		return kvm_arch_assign_device(dev);
+	group = kvm_vfio_file_iommu_group(file);
+	if (group)
+		return kvm_arch_assign_group(group);
+
+	return -ENODEV;
+}
+
+static void kvm_vfio_reclaim_file(struct file *file)
+{
+	struct device *dev;
+	struct iommu_group *group;
+
+	dev = kvm_vfio_file_get_device(file);
+	if (dev) {
+		kvm_arch_reclaim_device(dev);
+		return;
+	}
+	group = kvm_vfio_file_iommu_group(file);
+	if (group)
+		kvm_arch_reclaim_group(group);
+}
+
 static int kvm_vfio_file_add(struct kvm_device *dev, unsigned int fd)
 {
 	struct kvm_vfio *kv = dev->private;
@@ -172,6 +218,10 @@ static int kvm_vfio_file_add(struct kvm_device *dev, unsigned int fd)
 		goto out_unlock;
 	}
 
+	ret = kvm_vfio_assign_file(filp);
+	if (ret)
+		goto out_unlock;
+
 	kvf->file = get_file(filp);
 	list_add_tail(&kvf->node, &kv->file_list);
 
@@ -205,6 +255,7 @@ static int kvm_vfio_file_del(struct kvm_device *dev, unsigned int fd)
 		if (kvf->file != fd_file(f))
 			continue;
 
+		kvm_vfio_reclaim_file(kvf->file);
 		list_del(&kvf->node);
 		kvm_arch_end_assignment(dev->kvm);
 #ifdef CONFIG_SPAPR_TCE_IOMMU
@@ -338,6 +389,7 @@ static void kvm_vfio_release(struct kvm_device *dev)
 #ifdef CONFIG_SPAPR_TCE_IOMMU
 		kvm_spapr_tce_release_vfio_group(dev->kvm, kvf);
 #endif
+		kvm_vfio_reclaim_file(kvf->file);
 		kvm_vfio_file_set_kvm(kvf->file, NULL);
 		fput(kvf->file);
 		list_del(&kvf->node);