diff --git a/drivers/vfio/platform/vfio_platform_common.c b/drivers/vfio/platform/vfio_platform_common.c
index e65b142d3422..993b2f932f80 100644
--- a/drivers/vfio/platform/vfio_platform_common.c
+++ b/drivers/vfio/platform/vfio_platform_common.c
@@ -561,7 +561,7 @@ int vfio_platform_probe_common(struct vfio_platform_device *vdev,
 
 	vdev->device = dev;
 
-	group = iommu_group_get(dev);
+	group = vfio_iommu_group_get(dev);
 	if (!group) {
 		pr_err("VFIO: No IOMMU group for device %s\n", vdev->name);
 		return -EINVAL;
@@ -569,7 +569,7 @@ int vfio_platform_probe_common(struct vfio_platform_device *vdev,
 
 	ret = vfio_add_group_dev(dev, &vfio_platform_ops, vdev);
 	if (ret) {
-		iommu_group_put(group);
+		vfio_iommu_group_put(group, dev);
 		return ret;
 	}
 
@@ -589,7 +589,7 @@ struct vfio_platform_device *vfio_platform_remove_common(struct device *dev)
 
 	if (vdev) {
 		vfio_platform_put_reset(vdev);
-		iommu_group_put(dev->iommu_group);
+		vfio_iommu_group_put(dev->iommu_group, dev);
 	}
 
 	return vdev;