#include <linux/types.h>
 
-/* PTE_PFN_MASK extracts the PFN from a (pte|pmd|pud|pgd)val_t */
+/* Extracts the PFN from a (pte|pmd|pud|pgd)val_t of a 4KB page */
 #define PTE_PFN_MASK           ((pteval_t)PHYSICAL_PAGE_MASK)
 
-/* PTE_FLAGS_MASK extracts the flags from a (pte|pmd|pud|pgd)val_t */
+/* Extracts the flags from a (pte|pmd|pud|pgd)val_t of a 4KB page */
 #define PTE_FLAGS_MASK         (~PTE_PFN_MASK)
 
 typedef struct pgprot { pgprotval_t pgprot; } pgprot_t;
 }
 #endif
 
+static inline pudval_t pud_pfn_mask(pud_t pud)
+{
+       if (native_pud_val(pud) & _PAGE_PSE)
+               return PUD_PAGE_MASK & PHYSICAL_PAGE_MASK;
+       else
+               return PTE_PFN_MASK;
+}
+
+static inline pudval_t pud_flags_mask(pud_t pud)
+{
+       if (native_pud_val(pud) & _PAGE_PSE)
+               return ~(PUD_PAGE_MASK & (pudval_t)PHYSICAL_PAGE_MASK);
+       else
+               return ~PTE_PFN_MASK;
+}
+
 static inline pudval_t pud_flags(pud_t pud)
 {
        return native_pud_val(pud) & PTE_FLAGS_MASK;
 }
 
+static inline pmdval_t pmd_pfn_mask(pmd_t pmd)
+{
+       if (native_pmd_val(pmd) & _PAGE_PSE)
+               return PMD_PAGE_MASK & PHYSICAL_PAGE_MASK;
+       else
+               return PTE_PFN_MASK;
+}
+
+static inline pmdval_t pmd_flags_mask(pmd_t pmd)
+{
+       if (native_pmd_val(pmd) & _PAGE_PSE)
+               return ~(PMD_PAGE_MASK & (pmdval_t)PHYSICAL_PAGE_MASK);
+       else
+               return ~PTE_PFN_MASK;
+}
+
 static inline pmdval_t pmd_flags(pmd_t pmd)
 {
        return native_pmd_val(pmd) & PTE_FLAGS_MASK;