ret);
 }
 
+/* Return largest page shift based on "IO Page Sizes" output of ibm,query-pe-dma-window. */
+static int iommu_get_page_shift(u32 query_page_size)
+{
+       /* Supported IO page-sizes according to LoPAR */
+       const int shift[] = {
+               __builtin_ctzll(SZ_4K),   __builtin_ctzll(SZ_64K), __builtin_ctzll(SZ_16M),
+               __builtin_ctzll(SZ_32M),  __builtin_ctzll(SZ_64M), __builtin_ctzll(SZ_128M),
+               __builtin_ctzll(SZ_256M), __builtin_ctzll(SZ_16G)
+       };
+
+       int i = ARRAY_SIZE(shift) - 1;
+
+       /*
+        * On LoPAR, ibm,query-pe-dma-window outputs "IO Page Sizes" using a bit field:
+        * - bit 31 means 4k pages are supported,
+        * - bit 30 means 64k pages are supported, and so on.
+        * Larger pagesizes map more memory with the same amount of TCEs, so start probing them.
+        */
+       for (; i >= 0 ; i--) {
+               if (query_page_size & (1 << i))
+                       return shift[i];
+       }
+
+       /* No valid page size found. */
+       return 0;
+}
+
 /*
  * If the PE supports dynamic dma windows, and there is space for a table
  * that can map all pages in a linear offset, then setup such a table,
                        goto out_failed;
                }
        }
-       if (query.page_size & 4) {
-               page_shift = 24; /* 16MB */
-       } else if (query.page_size & 2) {
-               page_shift = 16; /* 64kB */
-       } else if (query.page_size & 1) {
-               page_shift = 12; /* 4kB */
-       } else {
+
+       page_shift = iommu_get_page_shift(query.page_size);
+       if (!page_shift) {
                dev_dbg(&dev->dev, "no supported direct page size in mask %x",
                          query.page_size);
                goto out_failed;