diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index 1d0a70c85333aa..ad124959a26a63 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -439,6 +439,7 @@
 					      translation */
 #define PD_IOMMUV2_MASK		(1UL << 3) /* domain has gcr3 table */
 #define PD_GIOV_MASK		(1UL << 4) /* domain enable GIOV support */
+#define PD_VFIO_MASK		(1UL << 5) /* domain enable VFIO support */
 
 extern bool amd_iommu_dump;
 #define DUMP_printk(format, arg...)				\
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index d3b39d0416fa32..85074f3c4e0e44 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -2057,6 +2057,8 @@ static struct protection_domain *protection_domain_alloc(unsigned int type)
 		mode = PAGE_MODE_NONE;
 	} else if (type == IOMMU_DOMAIN_UNMANAGED) {
 		pgtable = AMD_IOMMU_V1;
+		/* Mark unmanaged domain for VFIO */
+		domain->flags |= PD_VFIO_MASK;
 	}
 
 	switch (pgtable) {
@@ -2131,6 +2133,7 @@ static void amd_iommu_detach_device(struct iommu_domain *dom,
 				    struct device *dev)
 {
 	struct iommu_dev_data *dev_data = dev_iommu_priv_get(dev);
+	struct protection_domain *domain = to_pdomain(dom);
 	struct amd_iommu *iommu;
 
 	if (!check_device(dev))
@@ -2145,7 +2148,7 @@ static void amd_iommu_detach_device(struct iommu_domain *dom,
 
 #ifdef CONFIG_IRQ_REMAP
 	if (AMD_IOMMU_GUEST_IR_VAPIC(amd_iommu_guest_ir) &&
-	    (dom->type == IOMMU_DOMAIN_UNMANAGED))
+	    (domain->flags & PD_VFIO_MASK))
 		dev_data->use_vapic = 0;
 #endif
 
@@ -2177,7 +2180,7 @@ static int amd_iommu_attach_device(struct iommu_domain *dom,
 
 #ifdef CONFIG_IRQ_REMAP
 	if (AMD_IOMMU_GUEST_IR_VAPIC(amd_iommu_guest_ir)) {
-		if (dom->type == IOMMU_DOMAIN_UNMANAGED)
+		if (domain->flags & PD_VFIO_MASK)
 			dev_data->use_vapic = 1;
 		else
 			dev_data->use_vapic = 0;