WIP: ANDROID: iommu/io-pgtable-arm: Coalesce beyond last level

Coalesce levels before the last level possible, this runs post table
to make sure any last level coalescing ran first so we only coalesce
blocks to a table.

Change-Id: Id9d6c6208b2179657515eb90e57dea927829e5ec
Signed-off-by: Mostafa Saleh <smostafa@google.com>
diff --git a/drivers/iommu/arm/arm-smmu-v3/pkvm/io-pgtable-arm.c b/drivers/iommu/arm/arm-smmu-v3/pkvm/io-pgtable-arm.c
index 5f49bb8..d604ff9 100644
--- a/drivers/iommu/arm/arm-smmu-v3/pkvm/io-pgtable-arm.c
+++ b/drivers/iommu/arm/arm-smmu-v3/pkvm/io-pgtable-arm.c
@@ -246,3 +246,32 @@ bool arm_lpae_use_block_mapping(struct arm_lpae_io_pgtable *data,
 
 	return false;
 }
+
+void arm_lpae_post_table_walk(struct arm_lpae_io_pgtable *data,
+			      unsigned long iova, size_t pgsize,
+			      size_t pgcount, arm_lpae_iopte prot,
+			      int level, arm_lpae_iopte *ptep)
+{
+	size_t block_size = ARM_LPAE_BLOCK_SIZE(level, data);
+	arm_lpae_iopte pte = *ptep;
+	int i = 0;
+	arm_lpae_iopte *cptep = iopte_deref(pte, data);
+
+	if (!data->idmapped)
+		return;
+
+	/* Already done. */
+	if (level >= ARM_LPAE_MAX_LEVELS - 2)
+		return;
+
+	for (i = 0 ; i < ARM_LPAE_PTES_PER_TABLE(data); ++i) {
+		arm_lpae_iopte pte = cptep[i];
+		if (!iopte_leaf(pte, level + 1, data->iop.fmt) || arm_lpae_iopte_is_mmio(data, pte))
+			return;
+	}
+
+	iova &= ~(block_size - 1);
+	WARN_ON(arm_lpae_init_pte(data, iova, iova, prot, level, 1, ptep));
+
+	return;
+}
diff --git a/drivers/iommu/io-pgtable-arm-common.c b/drivers/iommu/io-pgtable-arm-common.c
index b692c21..775e746 100644
--- a/drivers/iommu/io-pgtable-arm-common.c
+++ b/drivers/iommu/io-pgtable-arm-common.c
@@ -74,10 +74,10 @@ static void __arm_lpae_init_pte(struct arm_lpae_io_pgtable *data,
 		__arm_lpae_sync_pte(ptep, num_entries, cfg);
 }
 
-static int arm_lpae_init_pte(struct arm_lpae_io_pgtable *data,
-			     unsigned long iova, phys_addr_t paddr,
-			     arm_lpae_iopte prot, int lvl, int num_entries,
-			     arm_lpae_iopte *ptep)
+int arm_lpae_init_pte(struct arm_lpae_io_pgtable *data,
+		      unsigned long iova, phys_addr_t paddr,
+		      arm_lpae_iopte prot, int lvl, int num_entries,
+		      arm_lpae_iopte *ptep)
 {
 	int i, mapped_entries = 0;
 	size_t sz = ARM_LPAE_BLOCK_SIZE(lvl, data);
@@ -215,8 +215,12 @@ int __arm_lpae_map(struct arm_lpae_io_pgtable *data, unsigned long iova,
 	}
 
 	/* Rinse, repeat */
-	return __arm_lpae_map(data, iova, paddr, size, pgcount, prot, lvl + 1,
-			      cptep, gfp, mapped);
+	ret = __arm_lpae_map(data, iova, paddr, size, pgcount, prot, lvl + 1,
+			     cptep, gfp, mapped);
+
+	/* Trigger coalescing logic if supported. */
+	arm_lpae_post_table_walk(data, iova, size, pgcount, prot, lvl, ptep);
+	return ret;
 }
 
 static arm_lpae_iopte arm_lpae_prot_to_pte(struct arm_lpae_io_pgtable *data,
diff --git a/drivers/iommu/io-pgtable-arm.c b/drivers/iommu/io-pgtable-arm.c
index 4cd7f64c..525a082 100644
--- a/drivers/iommu/io-pgtable-arm.c
+++ b/drivers/iommu/io-pgtable-arm.c
@@ -263,6 +263,14 @@ bool arm_lpae_use_block_mapping(struct arm_lpae_io_pgtable *data,
 	return pgsize == block_size;
 }
 
+void arm_lpae_post_table_walk(struct arm_lpae_io_pgtable *data,
+			      unsigned long iova, size_t pgsize,
+			      size_t pgcount, arm_lpae_iopte prot,
+			      int level, arm_lpae_iopte *ptep)
+{
+	return;
+}
+
 int arm_lpae_mapping_exists(struct arm_lpae_io_pgtable *data)
 {
 	/* We require an unmap first */
diff --git a/include/linux/io-pgtable-arm.h b/include/linux/io-pgtable-arm.h
index 25e252d..77507bf 100644
--- a/include/linux/io-pgtable-arm.h
+++ b/include/linux/io-pgtable-arm.h
@@ -217,7 +217,15 @@ bool arm_lpae_use_block_mapping(struct arm_lpae_io_pgtable *data,
 				unsigned long iova, size_t pgsize,
 				size_t *pgcount, arm_lpae_iopte prot,
 				int level, arm_lpae_iopte *ptep);
+void arm_lpae_post_table_walk(struct arm_lpae_io_pgtable *data,
+			      unsigned long iova, size_t pgsize,
+			      size_t pgcount, arm_lpae_iopte prot,
+			      int level, arm_lpae_iopte *ptep);
 int arm_lpae_mapping_exists(struct arm_lpae_io_pgtable *data);
 void arm_lpae_mapping_missing(struct arm_lpae_io_pgtable *data);
 
+int arm_lpae_init_pte(struct arm_lpae_io_pgtable *data,
+		      unsigned long iova, phys_addr_t paddr,
+		      arm_lpae_iopte prot, int lvl, int num_entries,
+		      arm_lpae_iopte *ptep);
 #endif /* IO_PGTABLE_H_ */