return (paddr | (paddr << (48 - 12))) & (ARM_LPAE_PTE_ADDR_MASK << 4);
 }
 
+/*
+ * Convert an index returned by ARM_LPAE_PGD_IDX(), which can point into
+ * a concatenated PGD, into the maximum number of entries that can be
+ * mapped in the same table page.
+ */
+static inline int arm_lpae_max_entries(int i, struct arm_lpae_io_pgtable *data)
+{
+       int ptes_per_table = ARM_LPAE_PTES_PER_TABLE(data);
+
+       return ptes_per_table - (i & (ptes_per_table - 1));
+}
+
 static bool selftest_running = false;
 
 static dma_addr_t __arm_lpae_dma_addr(void *pages)
 
        /* If we can install a leaf entry at this level, then do so */
        if (size == block_size) {
-               max_entries = ARM_LPAE_PTES_PER_TABLE(data) - map_idx_start;
+               max_entries = arm_lpae_max_entries(map_idx_start, data);
                num_entries = min_t(int, pgcount, max_entries);
                ret = arm_lpae_init_pte(data, iova, paddr, prot, lvl, num_entries, ptep);
                if (!ret)
 
        if (size == split_sz) {
                unmap_idx_start = ARM_LPAE_LVL_IDX(iova, lvl, data);
-               max_entries = ptes_per_table - unmap_idx_start;
+               max_entries = arm_lpae_max_entries(unmap_idx_start, data);
                num_entries = min_t(int, pgcount, max_entries);
        }
 
 
        /* If the size matches this level, we're in the right place */
        if (size == ARM_LPAE_BLOCK_SIZE(lvl, data)) {
-               max_entries = ARM_LPAE_PTES_PER_TABLE(data) - unmap_idx_start;
+               max_entries = arm_lpae_max_entries(unmap_idx_start, data);
                num_entries = min_t(int, pgcount, max_entries);
 
                /* Find and handle non-leaf entries */