return ptes_per_table - (i & (ptes_per_table - 1));
 }
 
+/*
+ * Check if concatenated PGDs are mandatory according to Arm DDI0487 (K.a)
+ * 1) R_DXBSH: For 16KB, and 48-bit input size, use level 1 instead of 0.
+ * 2) R_SRKBC: After de-ciphering the table for PA size and valid initial lookup
+ *   a) 40 bits PA size with 4K: use level 1 instead of level 0 (2 tables for ias = oas)
+ *   b) 40 bits PA size with 16K: use level 2 instead of level 1 (16 tables for ias = oas)
+ *   c) 42 bits PA size with 4K: use level 1 instead of level 0 (8 tables for ias = oas)
+ *   d) 48 bits PA size with 16K: use level 1 instead of level 0 (2 tables for ias = oas)
+ */
+static inline bool arm_lpae_concat_mandatory(struct arm_lpae_io_pgtable *data)
+{
+       unsigned int ias = data->iop.cfg.ias;
+       unsigned int oas = data->iop.cfg.oas;
+
+       /* Covers 1  and 2.d */
+       if ((ARM_LPAE_GRANULE(data) == SZ_16K) && (data->start_level == 0))
+               return (oas == 48) || (ias == 48);
+
+       /* Covers 2.a and 2.c */
+       if ((ARM_LPAE_GRANULE(data) == SZ_4K) && (data->start_level == 0))
+               return (oas == 40) || (oas == 42);
+
+       /* Case 2.b */
+       return (ARM_LPAE_GRANULE(data) == SZ_16K) &&
+              (data->start_level == 1) && (oas == 40);
+}
+
 static bool selftest_running = false;
 
 static dma_addr_t __arm_lpae_dma_addr(void *pages)
        if (!data)
                return NULL;
 
-       /*
-        * Concatenate PGDs at level 1 if possible in order to reduce
-        * the depth of the stage-2 walk.
-        */
-       if (data->start_level == 0) {
-               unsigned long pgd_pages;
-
-               pgd_pages = ARM_LPAE_PGD_SIZE(data) / sizeof(arm_lpae_iopte);
-               if (pgd_pages <= ARM_LPAE_S2_MAX_CONCAT_PAGES) {
-                       data->pgd_bits += data->bits_per_level;
-                       data->start_level++;
-               }
+       if (arm_lpae_concat_mandatory(data))  {
+               if (WARN_ON((ARM_LPAE_PGD_SIZE(data) / sizeof(arm_lpae_iopte)) >
+                           ARM_LPAE_S2_MAX_CONCAT_PAGES))
+                       return NULL;
+               data->pgd_bits += data->bits_per_level;
+               data->start_level++;
        }
 
        /* VTCR */