iommu/amd: Add function to get IOMMUv2 domain for pdev

The AMD IOMMUv2 driver needs to get the IOMMUv2 domain
associated with a particular device. This patch adds a
function to get this information.

Signed-off-by: Joerg Roedel <joerg.roedel@amd.com>
diff --git a/drivers/iommu/amd_iommu.c b/drivers/iommu/amd_iommu.c
index 9a7e64b..71773d0 100644
--- a/drivers/iommu/amd_iommu.c
+++ b/drivers/iommu/amd_iommu.c
@@ -2958,6 +2958,8 @@
 	if (!domain->pt_root)
 		goto out_free;
 
+	domain->iommu_domain = dom;
+
 	dom->priv = domain;
 
 	return 0;
@@ -3463,3 +3465,19 @@
 	return iommu_queue_command(iommu, &cmd);
 }
 EXPORT_SYMBOL(amd_iommu_complete_ppr);
+
+struct iommu_domain *amd_iommu_get_v2_domain(struct pci_dev *pdev)
+{
+	struct protection_domain *domain;
+
+	domain = get_domain(&pdev->dev);
+	if (IS_ERR(domain))
+		return NULL;
+
+	/* Only return IOMMUv2 domains */
+	if (!(domain->flags & PD_IOMMUV2_MASK))
+		return NULL;
+
+	return domain->iommu_domain;
+}
+EXPORT_SYMBOL(amd_iommu_get_v2_domain);