struct panfrost_mmu *panfrost_mmu_ctx_create(struct panfrost_device *pfdev)
 {
+       u32 va_bits = GPU_MMU_FEATURES_VA_BITS(pfdev->features.mmu_features);
+       u32 pa_bits = GPU_MMU_FEATURES_PA_BITS(pfdev->features.mmu_features);
        struct panfrost_mmu *mmu;
 
        mmu = kzalloc(sizeof(*mmu), GFP_KERNEL);
 
        mmu->pgtbl_cfg = (struct io_pgtable_cfg) {
                .pgsize_bitmap  = SZ_4K | SZ_2M,
-               .ias            = FIELD_GET(0xff, pfdev->features.mmu_features),
-               .oas            = FIELD_GET(0xff00, pfdev->features.mmu_features),
+               .ias            = va_bits,
+               .oas            = pa_bits,
                .coherent_walk  = pfdev->coherent,
                .tlb            = &mmu_tlb_ops,
                .iommu_dev      = pfdev->dev,
 
 #define   GROUPS_L2_COHERENT           BIT(0)  /* Cores groups are l2 coherent */
 
 #define GPU_MMU_FEATURES               0x014   /* (RO) MMU features */
+#define  GPU_MMU_FEATURES_VA_BITS(x)   ((x) & GENMASK(7, 0))
+#define  GPU_MMU_FEATURES_PA_BITS(x)   (((x) >> 8) & GENMASK(7, 0))
 #define GPU_AS_PRESENT                 0x018   /* (RO) Address space slots present */
 #define GPU_JS_PRESENT                 0x01C   /* (RO) Job slots present */