diff --git a/src/arch/aarch64/inc/hf/arch/mm.h b/src/arch/aarch64/inc/hf/arch/mm.h
index eaf7c4c..980f386 100644
--- a/src/arch/aarch64/inc/hf/arch/mm.h
+++ b/src/arch/aarch64/inc/hf/arch/mm.h
@@ -10,112 +10,121 @@
 
 #define PAGE_LEVEL_BITS 9
 
+#define HF_ARCH_AARCH64_MM_PTE_ATTR_MASK \
+	(((UINT64_C(1) << PAGE_BITS) - 1) | ~((UINT64_C(1) << 48) - 1))
+
+#define HF_ARCH_AARCH64_MM_GET_PTE_ATTRS(v) \
+	((v)&HF_ARCH_AARCH64_MM_PTE_ATTR_MASK)
+
+#define HF_ARCH_AARCH64_MM_CLEAR_PTE_ATTRS(v) \
+	((v) & ~HF_ARCH_AARCH64_MM_PTE_ATTR_MASK)
+
+/**
+ * Returns the encoding of a page table entry that isn't present.
+ */
+static inline pte_t arch_mm_absent_pte(int level)
+{
+	return 0;
+}
+
 /**
  * Converts a physical address to a table PTE.
  *
  * The spec says that 'Table descriptors for stage 2 translations do not
  * include any attribute field', so we don't take any attributes as arguments.
  */
-static inline pte_t arch_mm_pa_to_table_pte(paddr_t pa)
+static inline pte_t arch_mm_table_pte(int level, paddr_t pa)
 {
+	/* This is the same for all levels on aarch64. */
+	(void)level;
 	return pa_addr(pa) | 0x3;
 }
 
 /**
  * Converts a physical address to a block PTE.
+ *
+ * The level must allow block entries.
  */
-static inline pte_t arch_mm_pa_to_block_pte(paddr_t pa, uint64_t attrs)
+static inline pte_t arch_mm_block_pte(int level, paddr_t pa, uint64_t attrs)
 {
-	return pa_addr(pa) | attrs;
-}
-
-/**
- * Converts a physical address to a page PTE.
- */
-static inline pte_t arch_mm_pa_to_page_pte(paddr_t pa, uint64_t attrs)
-{
-	return pa_addr(pa) | attrs | ((attrs & 1) << 1);
-}
-
-/**
- * Converts a block PTE to a page PTE.
- */
-static inline pte_t arch_mm_block_to_page_pte(pte_t pte)
-{
-	return pte | 2;
+	pte_t pte = pa_addr(pa) | attrs;
+	if (level == 0) {
+		pte |= 0x2;
+	}
+	return pte;
 }
 
 /**
  * Specifies whether block mappings are acceptable at the given level.
+ *
+ * Level 0 must allow block entries.
  */
 static inline bool arch_mm_is_block_allowed(int level)
 {
-	return level == 1 || level == 2;
-}
-
-/**
- * Returns the encoding of a page table entry that isn't present.
- */
-static inline pte_t arch_mm_absent_pte(void)
-{
-	return 0;
+	return level <= 2;
 }
 
 /**
  * Determines if the given pte is present, i.e., if it points to another table,
  * to a page, or a block of pages.
  */
-static inline bool arch_mm_pte_is_present(pte_t pte)
+static inline bool arch_mm_pte_is_present(pte_t pte, int level)
 {
-	return (pte & 1) != 0;
+	return (pte & 0x1) != 0;
 }
 
 /**
  * Determines if the given pte references another table.
  */
-static inline bool arch_mm_pte_is_table(pte_t pte)
+static inline bool arch_mm_pte_is_table(pte_t pte, int level)
 {
-	return (pte & 3) == 3;
+	return level != 0 && (pte & 0x3) == 0x3;
 }
 
 /**
  * Determines if the given pte references a block of pages.
  */
-static inline bool arch_mm_pte_is_block(pte_t pte)
+static inline bool arch_mm_pte_is_block(pte_t pte, int level)
 {
-	return (pte & 3) == 1;
+	return arch_mm_is_block_allowed(level) &&
+	       (pte & 0x3) == (level == 0 ? 0x3 : 0x1);
 }
 
-#define CLEAR_PTE_ATTRS(v) \
-	((v) & ~((1ull << PAGE_BITS) - 1) & ((1ull << 48) - 1))
-
 /**
  * Clears the given physical address, i.e., sets the ignored bits (from a page
  * table perspective) to zero.
  */
 static inline paddr_t arch_mm_clear_pa(paddr_t pa)
 {
-	return pa_init(CLEAR_PTE_ATTRS(pa_addr(pa)));
+	return pa_init(HF_ARCH_AARCH64_MM_CLEAR_PTE_ATTRS(pa_addr(pa)));
 }
 
 /**
- * Extracts the physical address from a page table entry.
+ * Extracts the physical address of the block referred to by the given page
+ * table entry.
  */
-static inline paddr_t arch_mm_pte_to_paddr(pte_t pte)
+static inline paddr_t arch_mm_block_from_pte(pte_t pte)
 {
-	return pa_init(CLEAR_PTE_ATTRS(pte));
+	return pa_init(HF_ARCH_AARCH64_MM_CLEAR_PTE_ATTRS(pte));
 }
 
 /**
  * Extracts the physical address of the page table referred to by the given page
  * table entry.
  */
-static inline paddr_t arch_mm_pte_to_table(pte_t pte)
+static inline paddr_t arch_mm_table_from_pte(pte_t pte)
 {
-	return pa_init(CLEAR_PTE_ATTRS(pte));
+	return pa_init(HF_ARCH_AARCH64_MM_CLEAR_PTE_ATTRS(pte));
 }
 
-#undef CLEAR_PTE_ATTRS
+/**
+ * Extracts the architecture specific attributes applies to the given page table
+ * entry.
+ */
+static inline uint64_t arch_mm_pte_attrs(pte_t pte)
+{
+	return HF_ARCH_AARCH64_MM_GET_PTE_ATTRS(pte);
+}
 
 /**
  * Invalidates stage-1 TLB entries referring to the given virtual address range.
@@ -132,7 +141,7 @@
 
 	__asm__ volatile("dsb ishst");
 
-	for (it = begin; it < end; it += (1ull << (PAGE_BITS - 12))) {
+	for (it = begin; it < end; it += (UINT64_C(1) << (PAGE_BITS - 12))) {
 		__asm__("tlbi vae2is, %0" : : "r"(it));
 	}
 
@@ -157,7 +166,7 @@
 
 	__asm__ volatile("dsb ishst");
 
-	for (it = begin; it < end; it += (1ull << (PAGE_BITS - 12))) {
+	for (it = begin; it < end; it += (UINT64_C(1) << (PAGE_BITS - 12))) {
 		__asm__("tlbi ipas2e1, %0" : : "r"(it));
 	}
 
diff --git a/src/arch/aarch64/mm.c b/src/arch/aarch64/mm.c
index 212a56e..df4ed2d 100644
--- a/src/arch/aarch64/mm.c
+++ b/src/arch/aarch64/mm.c
@@ -9,54 +9,54 @@
 /* Keep macro alignment */
 /* clang-format off */
 
-#define NON_SHAREABLE   0ull
-#define OUTER_SHAREABLE 2ull
-#define INNER_SHAREABLE 3ull
+#define NON_SHAREABLE   UINT64_C(0)
+#define OUTER_SHAREABLE UINT64_C(2)
+#define INNER_SHAREABLE UINT64_C(3)
 
-#define STAGE1_XN          (1ull << 54)
-#define STAGE1_CONTIGUOUS  (1ull << 52)
-#define STAGE1_DBM         (1ull << 51)
-#define STAGE1_NG          (1ull << 11)
-#define STAGE1_AF          (1ull << 10)
+#define STAGE1_XN          (UINT64_C(1) << 54)
+#define STAGE1_CONTIGUOUS  (UINT64_C(1) << 52)
+#define STAGE1_DBM         (UINT64_C(1) << 51)
+#define STAGE1_NG          (UINT64_C(1) << 11)
+#define STAGE1_AF          (UINT64_C(1) << 10)
 #define STAGE1_SH(x)       ((x) << 8)
 #define STAGE1_AP(x)       ((x) << 6)
-#define STAGE1_NS          (1ull << 5)
+#define STAGE1_NS          (UINT64_C(1) << 5)
 #define STAGE1_ATTRINDX(x) ((x) << 2)
 
-#define STAGE1_READONLY  2ull
-#define STAGE1_READWRITE 0ull
+#define STAGE1_READONLY  UINT64_C(2)
+#define STAGE1_READWRITE UINT64_C(0)
 
-#define STAGE1_DEVICEINDX 0ull
-#define STAGE1_NORMALINDX 1ull
+#define STAGE1_DEVICEINDX UINT64_C(0)
+#define STAGE1_NORMALINDX UINT64_C(1)
 
 #define STAGE2_XN(x)      ((x) << 53)
-#define STAGE2_CONTIGUOUS (1ull << 52)
-#define STAGE2_DBM        (1ull << 51)
-#define STAGE2_AF         (1ull << 10)
+#define STAGE2_CONTIGUOUS (UINT64_C(1) << 52)
+#define STAGE2_DBM        (UINT64_C(1) << 51)
+#define STAGE2_AF         (UINT64_C(1) << 10)
 #define STAGE2_SH(x)      ((x) << 8)
 #define STAGE2_S2AP(x)    ((x) << 6)
 #define STAGE2_MEMATTR(x) ((x) << 2)
 
-#define STAGE2_EXECUTE_ALL  0ull
-#define STAGE2_EXECUTE_EL0  1ull
-#define STAGE2_EXECUTE_NONE 2ull
-#define STAGE2_EXECUTE_EL1  3ull
+#define STAGE2_EXECUTE_ALL  UINT64_C(0)
+#define STAGE2_EXECUTE_EL0  UINT64_C(1)
+#define STAGE2_EXECUTE_NONE UINT64_C(2)
+#define STAGE2_EXECUTE_EL1  UINT64_C(3)
 
 /* The following are stage-2 memory attributes for normal memory. */
-#define STAGE2_NONCACHEABLE 1ull
-#define STAGE2_WRITETHROUGH 2ull
-#define STAGE2_WRITEBACK    3ull
+#define STAGE2_NONCACHEABLE UINT64_C(1)
+#define STAGE2_WRITETHROUGH UINT64_C(2)
+#define STAGE2_WRITEBACK    UINT64_C(3)
 
 #define STAGE2_MEMATTR_NORMAL(outer, inner) ((((outer) << 2) | (inner)) << 2)
 
 /* The following stage-2 memory attributes for device memory. */
-#define STAGE2_MEMATTR_DEVICE_nGnRnE (0ull << 2)
-#define STAGE2_MEMATTR_DEVICE_nGnRE  (1ull << 2)
-#define STAGE2_MEMATTR_DEVICE_nGRE   (2ull << 2)
-#define STAGE2_MEMATTR_DEVICE_GRE    (3ull << 2)
+#define STAGE2_MEMATTR_DEVICE_nGnRnE (UINT64_C(0) << 2)
+#define STAGE2_MEMATTR_DEVICE_nGnRE  (UINT64_C(1) << 2)
+#define STAGE2_MEMATTR_DEVICE_nGRE   (UINT64_C(2) << 2)
+#define STAGE2_MEMATTR_DEVICE_GRE    (UINT64_C(3) << 2)
 
-#define STAGE2_ACCESS_READ  1ull
-#define STAGE2_ACCESS_WRITE 2ull
+#define STAGE2_ACCESS_READ  UINT64_C(1)
+#define STAGE2_ACCESS_WRITE UINT64_C(2)
 
 /* clang-format on */
 
diff --git a/src/load.c b/src/load.c
index b55455a..cd81237 100644
--- a/src/load.c
+++ b/src/load.c
@@ -123,11 +123,12 @@
 
 		/* Map the 1TB of memory. */
 		/* TODO: We should do a whitelist rather than a blacklist. */
-		if (!mm_vm_identity_map(&primary_vm.ptable, pa_init(0),
-					pa_init(1024ull * 1024 * 1024 * 1024),
-					MM_MODE_R | MM_MODE_W | MM_MODE_X |
-						MM_MODE_NOINVALIDATE,
-					NULL)) {
+		if (!mm_vm_identity_map(
+			    &primary_vm.ptable, pa_init(0),
+			    pa_init(UINT64_C(1024) * 1024 * 1024 * 1024),
+			    MM_MODE_R | MM_MODE_W | MM_MODE_X |
+				    MM_MODE_NOINVALIDATE,
+			    NULL)) {
 			dlog("Unable to initialise memory for primary vm\n");
 			return false;
 		}
diff --git a/src/mm.c b/src/mm.c
index d3b6a46..776f7a6 100644
--- a/src/mm.c
+++ b/src/mm.c
@@ -58,7 +58,7 @@
  */
 static inline size_t mm_entry_size(int level)
 {
-	return 1ull << (PAGE_BITS + level * PAGE_LEVEL_BITS);
+	return UINT64_C(1) << (PAGE_BITS + level * PAGE_LEVEL_BITS);
 }
 
 /**
@@ -78,7 +78,7 @@
 static inline size_t mm_index(ptable_addr_t addr, int level)
 {
 	ptable_addr_t v = addr >> (PAGE_BITS + level * PAGE_LEVEL_BITS);
-	return v & ((1ull << PAGE_LEVEL_BITS) - 1);
+	return v & ((UINT64_C(1) << PAGE_LEVEL_BITS) - 1);
 }
 
 /**
@@ -96,8 +96,8 @@
 	size_t inc;
 
 	/* Just return pointer to table if it's already populated. */
-	if (arch_mm_pte_is_table(v)) {
-		return ptr_from_va(va_from_pa(arch_mm_pte_to_table(v)));
+	if (arch_mm_pte_is_table(v, level)) {
+		return ptr_from_va(va_from_pa(arch_mm_table_from_pte(v)));
 	}
 
 	/* Allocate a new table. */
@@ -109,16 +109,15 @@
 	}
 
 	/* Determine template for new pte and its increment. */
-	if (!arch_mm_pte_is_block(v)) {
-		inc = 0;
-		new_pte = arch_mm_absent_pte();
+	if (arch_mm_pte_is_block(v, level)) {
+		int level_below = level - 1;
+		inc = mm_entry_size(level_below);
+		new_pte = arch_mm_block_pte(level_below,
+					    arch_mm_block_from_pte(v),
+					    arch_mm_pte_attrs(v));
 	} else {
-		inc = mm_entry_size(level - 1);
-		if (level == 1) {
-			new_pte = arch_mm_block_to_page_pte(v);
-		} else {
-			new_pte = v;
-		}
+		inc = 0;
+		new_pte = arch_mm_absent_pte(level);
 	}
 
 	/* Initialise entries in the new table. */
@@ -132,7 +131,7 @@
 	 * update it.
 	 */
 	atomic_thread_fence(memory_order_release);
-	*pte = arch_mm_pa_to_table_pte(pa_init((uintpaddr_t)ntable));
+	*pte = arch_mm_table_pte(level, pa_init((uintpaddr_t)ntable));
 
 	return ntable;
 }
@@ -147,7 +146,7 @@
 	(void)level;
 	(void)sync;
 	/* TODO: Implement.
-	if (!arch_mm_pte_is_present(pte) || level < 1)
+	if (!arch_mm_pte_is_present(pte, level) || level < 1)
 		return;
 	*/
 }
@@ -176,16 +175,12 @@
 
 	/* Fill each entry in the table. */
 	while (begin < end) {
-		if (level == 0) {
-			if (commit) {
-				*pte = arch_mm_pa_to_page_pte(pa, attrs);
-			}
-		} else if ((end - begin) >= entry_size &&
-			   arch_mm_is_block_allowed(level) &&
-			   (begin & (entry_size - 1)) == 0) {
+		if ((end - begin) >= entry_size &&
+		    arch_mm_is_block_allowed(level) &&
+		    (begin & (entry_size - 1)) == 0) {
 			if (commit) {
 				pte_t v = *pte;
-				*pte = arch_mm_pa_to_block_pte(pa, attrs);
+				*pte = arch_mm_block_pte(level, pa, attrs);
 				/* TODO: Add barrier. How do we ensure this
 				 * isn't in use by another CPU? Send IPI? */
 				mm_free_page_pte(v, level, sync);
@@ -321,7 +316,7 @@
 	}
 
 	i = mm_index(addr, 0);
-	table[i] = arch_mm_pa_to_page_pte(pa, attrs);
+	table[i] = arch_mm_block_pte(0, pa, attrs);
 	return true;
 }
 
@@ -333,19 +328,16 @@
 {
 	uint64_t i;
 	for (i = 0; i < PAGE_SIZE / sizeof(pte_t); i++) {
-		if (!arch_mm_pte_is_present(table[i])) {
+		if (!arch_mm_pte_is_present(table[i], level)) {
 			continue;
 		}
 
 		dlog("%*s%x: %x\n", 4 * (max_level - level), "", i, table[i]);
-		if (!level) {
-			continue;
-		}
 
-		if (arch_mm_pte_is_table(table[i])) {
+		if (arch_mm_pte_is_table(table[i], level)) {
 			mm_dump_table_recursive(
 				ptr_from_va(va_from_pa(
-					arch_mm_pte_to_table(table[i]))),
+					arch_mm_table_from_pte(table[i]))),
 				level - 1, max_level);
 		}
 	}
@@ -403,20 +395,17 @@
 
 	pte = table[mm_index(addr, level)];
 
-	if (level == 0) {
-		return arch_mm_pte_is_present(pte);
-	}
-
-	if (arch_mm_is_block_allowed(level) && arch_mm_pte_is_block(pte)) {
+	if (arch_mm_pte_is_block(pte, level)) {
 		return true;
 	}
 
-	if (arch_mm_pte_is_table(pte)) {
+	if (arch_mm_pte_is_table(pte, level)) {
 		return mm_is_mapped_recursive(
-			ptr_from_va(va_from_pa(arch_mm_pte_to_table(pte))),
+			ptr_from_va(va_from_pa(arch_mm_table_from_pte(pte))),
 			addr, level - 1);
 	}
 
+	/* The entry is not present. */
 	return false;
 }
 
@@ -453,7 +442,7 @@
 	}
 
 	for (i = 0; i < PAGE_SIZE / sizeof(pte_t); i++) {
-		table[i] = arch_mm_absent_pte();
+		table[i] = arch_mm_absent_pte(arch_mm_max_level(mode));
 	}
 
 	/* TODO: halloc could return a virtual or physical address if mm not