From 53e073efe162e14d3e0ee99379af832513138c7a Mon Sep 17 00:00:00 2001
From: Mostafa Saleh <smostafa@google.com>
Date: Fri, 4 Oct 2024 14:03:32 +0000
Subject: [PATCH] ANDROID: iommu/pkvm-iommu: Support devices in the same iommu
 group

At the moment, the driver assumes each device belongs to a different
group.

To support multiple devices in the same group, an optional ID can
be passed in the second field of iommu-cells, where devices
belonging to the same group have the same ID.

Bug: 357781595
Bug: 348382247
Bug: 236685427
Change-Id: I3627f951c45d2a2d7826546041df572d3d75b292
Signed-off-by: Mostafa Saleh <smostafa@google.com>
---
 .../bindings/iommu/pkvm,pviommu.yaml          | 10 ++++-
 drivers/iommu/pkvm-pviommu.c                  | 40 +++++++++++++++++--
 2 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/Documentation/devicetree/bindings/iommu/pkvm,pviommu.yaml b/Documentation/devicetree/bindings/iommu/pkvm,pviommu.yaml
index a44c74e1eb3c..070a01aad6c6 100644
--- a/Documentation/devicetree/bindings/iommu/pkvm,pviommu.yaml
+++ b/Documentation/devicetree/bindings/iommu/pkvm,pviommu.yaml
@@ -26,7 +26,15 @@ properties:
     maxItems: 1
 
   '#iommu-cells':
-    const: 1
+    enum: [ 1, 2 ]
+    description: |
+      See Documentation/devicetree/bindings/iommu/iommu.txt for details. With a
+      value of 1, each IOMMU specifier represents a distinct stream ID emitted
+      by that device into the relevant SMMU.
+
+      pvIOMMU would assume each device belongs to a different group, to describe
+      iommu groups, it is optionally passed in the second field of "iommus", where
+      devices having the same ID are part of the same group.
 
 required:
   - compatible
diff --git a/drivers/iommu/pkvm-pviommu.c b/drivers/iommu/pkvm-pviommu.c
index 478a2ddab268..5f9e1831ce50 100644
--- a/drivers/iommu/pkvm-pviommu.c
+++ b/drivers/iommu/pkvm-pviommu.c
@@ -9,6 +9,9 @@
 #include <linux/maple_tree.h>
 #include <linux/pci.h>
 #include <linux/platform_device.h>
+#include <linux/xarray.h>
+
+static DEFINE_XARRAY(pviommu_groups);
 
 struct pviommu_domain {
 	struct iommu_domain		domain;
@@ -339,15 +342,44 @@ static void pviommu_release_device(struct device *dev)
 
 static int pviommu_of_xlate(struct device *dev, const struct of_phandle_args *args)
 {
-	return iommu_fwspec_add_ids(dev, args->args, 1);
+	return iommu_fwspec_add_ids(dev, args->args, args->args_count);
+}
+
+static struct iommu_group *pviommu_group_alloc_get(struct device *dev, int group_id)
+{
+	struct iommu_group *group;
+
+	group = xa_load(&pviommu_groups, (unsigned long)group_id);
+	if (group)
+		return group;
+
+	group = iommu_group_alloc();
+	if (!IS_ERR(group))
+		return group;
+
+	if (WARN_ON(xa_insert(&pviommu_groups, (unsigned long)group_id, group, GFP_KERNEL)))
+		dev_err(dev,
+			"Failed to track group %d this will lead to multiple groups instead of one\n",
+			group_id);
+
+	return group;
 }
 
 static struct iommu_group *pviommu_device_group(struct device *dev)
 {
-	if (dev_is_pci(dev))
+	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
+
+	if (!fwspec)
+		return ERR_PTR(-ENODEV);
+
+	if (dev_is_pci(dev)) {
 		return pci_device_group(dev);
-	else
-		return generic_device_group(dev);
+	} else {
+		if (fwspec->num_ids == 1)
+			return generic_device_group(dev);
+		else
+			return pviommu_group_alloc_get(dev, fwspec->ids[1]);
+	}
 }
 
 static struct iommu_ops pviommu_ops = {