refactor(aarch64): Use aarch64-cpu for register access and some memory barriers

Definitions for DAIFSet were added recently, but are not yet in a
release, so the inline assembly that accesses it directly is marked with
a FIXME comment.

There are three differences I noticed while going through this:
- In the TCR setup, the `bfi` instruction was redundant since it was
  inserting the same 3 bits that were already shifted to the same
  position there earlier
- In TCR_EL1, 0 << 30 was being used to configure 4KiB granule size,
  which is incorrect. This is not even a valid value for TG1 (bits
  30-31). 0x10 is the correct value for 4KiB (the 0x01 for 16KiB was
  correct but unused)
- After enabling the MMU, the code performed a branch to flush the
  pipeline, which we can do more cleanly by inserting an instruction
  synchronization barrier

Author: Gelbpunkt

Cargo.lock

@@ -23,6 +23,7 @@ aarch64-cpu = "11"
 enum_dispatch = "0.3"
 fdt = { version = "0.1" }
 goblin = { version = "0.10", default-features = false, features = ["elf64"] }
+tock-registers = "0.10"
 volatile = { version = "0.6", features = ["derive"] }
 
 [target.'cfg(target_arch = "riscv64")'.dependencies]

Cargo.toml

@@ -2,48 +2,19 @@
 
 use core::arch::{asm, global_asm};
 
-use aarch64_cpu::registers::{SCTLR_EL1, Writeable};
+use aarch64_cpu::asm::{barrier, wfe};
+use aarch64_cpu::registers::{
+	CPACR_EL1, ID_AA64MMFR0_EL1, MAIR_EL1, MDSCR_EL1, ReadWriteable, Readable, SCTLR_EL1, TCR_EL1,
+	TPIDR_EL0, TPIDR_EL1, Writeable,
+};
 use log::info;
+use tock_registers::fields::{FieldValue, TryFromValue};
 
 const BOOT_CORE_ID: u64 = 0; // ID of CPU for booting on SMP systems - this might be board specific in the future
 
-/*
- * Memory types available.
- */
-#[allow(non_upper_case_globals)]
-const MT_DEVICE_nGnRnE: u64 = 0;
-#[allow(non_upper_case_globals)]
-const MT_DEVICE_nGnRE: u64 = 1;
-const MT_DEVICE_GRE: u64 = 2;
-const MT_NORMAL_NC: u64 = 3;
-const MT_NORMAL: u64 = 4;
-
-#[inline(always)]
-const fn mair(attr: u64, mt: u64) -> u64 {
-	attr << (mt * 8)
-}
-
-/*
- * TCR flags
- */
-const TCR_IRGN_WBWA: u64 = ((1) << 8) | ((1) << 24);
-const TCR_ORGN_WBWA: u64 = ((1) << 10) | ((1) << 26);
-const TCR_SHARED: u64 = ((3) << 12) | ((3) << 28);
-const TCR_TBI0: u64 = 1 << 37;
-const TCR_TBI1: u64 = 1 << 38;
-const TCR_ASID16: u64 = 1 << 36;
-const TCR_TG1_16K: u64 = 1 << 30;
-const TCR_TG1_4K: u64 = 0 << 30;
-const TCR_FLAGS: u64 = TCR_IRGN_WBWA | TCR_ORGN_WBWA | TCR_SHARED;
-
 /// Number of virtual address bits for 4KB page
 const VA_BITS: u64 = 48;
 
-#[inline(always)]
-const fn tcr_size(x: u64) -> u64 {
-	((64 - x) << 16) | (64 - x)
-}
-
 global_asm!(
 	include_str!("entry.s"),
 	start_rust = sym start_rust,
@@ -59,14 +30,14 @@ unsafe fn pre_init() -> ! {
 	info!("Enter startup code");
 
 	/* disable interrupts */
+	/* FIXME: Migrate to aarch64_cpu's DAIFSet definition once released */
 	unsafe {
 		asm!("msr daifset, 0b111", options(nostack));
 	}
 
 	/* reset thread id registers */
-	unsafe {
-		asm!("msr tpidr_el0, xzr", "msr tpidr_el1, xzr", options(nostack));
-	}
+	TPIDR_EL0.set(0);
+	TPIDR_EL1.set(0);
 
 	/*
 	 * Disable the MMU. We may have entered the kernel with it on and
@@ -75,87 +46,67 @@ unsafe fn pre_init() -> ! {
 	 * but in this case the code to find where we are running from
 	 * would have also failed.
 	 */
-	unsafe {
-		asm!("dsb sy",
-			"mrs x2, sctlr_el1",
-			"bic x2, x2, 0x1",
-			"msr sctlr_el1, x2",
-			"isb",
-			out("x2") _,
-			options(nostack),
-		);
-	}
+	barrier::dsb(barrier::SY);
+	SCTLR_EL1.modify(SCTLR_EL1::M::Disable);
+	barrier::isb(barrier::SY);
 
 	unsafe {
-		asm!("ic iallu", "tlbi vmalle1is", "dsb ish", options(nostack));
+		asm!("ic iallu", "tlbi vmalle1is", options(nostack));
 	}
+	barrier::dsb(barrier::ISH);
 
 	/*
 	 * Setup memory attribute type tables
-	 *
-	 * Memory regioin attributes for LPAE:
-	 *
-	 *   n = AttrIndx[2:0]
-	 *                      n       MAIR
-	 *   DEVICE_nGnRnE      000     00000000 (0x00)
-	 *   DEVICE_nGnRE       001     00000100 (0x04)
-	 *   DEVICE_GRE         010     00001100 (0x0c)
-	 *   NORMAL_NC          011     01000100 (0x44)
-	 *   NORMAL             100     11111111 (0xff)
 	 */
-	let mair_el1 = mair(0x00, MT_DEVICE_nGnRnE)
-		| mair(0x04, MT_DEVICE_nGnRE)
-		| mair(0x0c, MT_DEVICE_GRE)
-		| mair(0x44, MT_NORMAL_NC)
-		| mair(0xff, MT_NORMAL);
-	unsafe {
-		asm!("msr mair_el1, {}",
-			in(reg) mair_el1,
-			options(nostack),
-		);
-	}
+	MAIR_EL1.write(
+		MAIR_EL1::Attr0_Device::nonGathering_nonReordering_noEarlyWriteAck
+			+ MAIR_EL1::Attr1_Device::nonGathering_nonReordering_EarlyWriteAck
+			+ MAIR_EL1::Attr2_Device::Gathering_Reordering_EarlyWriteAck
+			+ MAIR_EL1::Attr3_Normal_Inner::NonCacheable
+			+ MAIR_EL1::Attr3_Normal_Outer::NonCacheable
+			+ MAIR_EL1::Attr4_Normal_Inner::WriteBack_NonTransient_ReadWriteAlloc
+			+ MAIR_EL1::Attr4_Normal_Outer::WriteBack_NonTransient_ReadWriteAlloc,
+	);
 
 	/*
 	 * Setup translation control register (TCR)
 	 */
 
 	// determine physical address size
-	unsafe {
-		asm!("mrs x0, id_aa64mmfr0_el1",
-			"and x0, x0, 0xF",
-			"lsl x0, x0, 32",
-			"orr x0, x0, {tcr_bits}",
-			"mrs x1, id_aa64mmfr0_el1",
-			"bfi x0, x1, #32, #3",
-			"msr tcr_el1, x0",
-			tcr_bits = in(reg) tcr_size(VA_BITS) | TCR_TG1_4K | TCR_FLAGS,
-			out("x0") _,
-			out("x1") _,
-		);
-	}
+	let Some(pa_range) = ID_AA64MMFR0_EL1
+		.read_as_enum::<ID_AA64MMFR0_EL1::PARange::Value>(ID_AA64MMFR0_EL1::PARange)
+	else {
+		panic!("Unknown physical address range")
+	};
+	let Some(ips) = TCR_EL1::IPS::Value::try_from_value(pa_range as u64) else {
+		panic!("Invalid physical address size")
+	};
+
+	TCR_EL1.write(
+		FieldValue::from(ips)
+			+ TCR_EL1::T1SZ.val(64 - VA_BITS)
+			+ TCR_EL1::T0SZ.val(64 - VA_BITS)
+			+ TCR_EL1::TG1::KiB_4
+			+ TCR_EL1::IRGN0::WriteBack_ReadAlloc_WriteAlloc_Cacheable
+			+ TCR_EL1::IRGN1::WriteBack_ReadAlloc_WriteAlloc_Cacheable
+			+ TCR_EL1::ORGN0::WriteBack_ReadAlloc_WriteAlloc_Cacheable
+			+ TCR_EL1::ORGN1::WriteBack_ReadAlloc_WriteAlloc_Cacheable
+			+ TCR_EL1::SH0::Inner
+			+ TCR_EL1::SH1::Inner,
+	);
 
 	/*
 	 * Enable FP/ASIMD in Architectural Feature Access Control Register,
 	 */
-	let bit_mask: u64 = 3 << 20;
-	unsafe {
-		asm!("msr cpacr_el1, {mask}",
-			mask = in(reg) bit_mask,
-			options(nostack),
-		);
-	}
+	CPACR_EL1.write(CPACR_EL1::FPEN::TrapNothing);
 
 	/*
 	 * Reset debug control register
 	 */
-	unsafe {
-		asm!("msr mdscr_el1, xzr", options(nostack));
-	}
+	MDSCR_EL1.set(0);
 
 	/* Memory barrier */
-	unsafe {
-		asm!("dsb sy", options(nostack));
-	}
+	barrier::dsb(barrier::SY);
 
 	/*
 	* Prepare system control register (SCTRL)
@@ -195,10 +146,8 @@ unsafe fn pre_init() -> ! {
 	}
 }
 
-pub unsafe fn wait_forever() -> ! {
+pub fn wait_forever() -> ! {
 	loop {
-		unsafe {
-			asm!("wfe");
-		}
+		wfe();
 	}
 }

src/arch/aarch64/entry.rs

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2025 The aarch64-rt Authors.
+ *
+ * This project is dual-licensed under Apache 2.0 and MIT terms.
+ * See LICENSE-APACHE and LICENSE-MIT for details.
+ */
+
+MEMORY
+{
+	image : ORIGIN = 0x40200000, LENGTH = 2M
+}

src/arch/aarch64/memory.ld

@@ -8,7 +8,8 @@ pub mod paging;
 use core::arch::asm;
 use core::ptr;
 
-use aarch64_cpu::asm::barrier::{NSH, SY, dmb, dsb, isb};
+use aarch64_cpu::asm::barrier::{self, NSH, SY, dmb, dsb, isb};
+use aarch64_cpu::registers::{ReadWriteable, SCTLR_EL1, TTBR0_EL1, TTBR1_EL1, Writeable};
 use align_address::Align;
 use fdt::Fdt;
 use goblin::elf::header::header64::{EI_DATA, ELFDATA2LSB, ELFMAG, Header, SELFMAG};
@@ -168,29 +169,14 @@ pub unsafe fn boot_kernel(kernel_info: LoadedKernel) -> ! {
 	CONSOLE.lock().get().set_stdout(0x1000);
 
 	// Load TTBRx
-	unsafe {
-		asm!(
-				"msr ttbr1_el1, xzr",
-				"msr ttbr0_el1, {}",
-				"dsb sy",
-				"isb",
-				in(reg) ptr::addr_of_mut!(l0_pgtable),
-				options(nostack),
-		)
-	};
+	TTBR1_EL1.set(0);
+	TTBR0_EL1.set(&raw mut l0_pgtable as u64);
+	barrier::dsb(barrier::SY);
+	barrier::isb(barrier::SY);
 
 	// Enable paging
-	unsafe {
-		asm!(
-				"mrs x0, sctlr_el1",
-				"orr x0, x0, #1",
-				"msr sctlr_el1, x0",
-				"bl 0f",
-				"0:",
-				out("x0") _,
-				options(nostack),
-		);
-	}
+	SCTLR_EL1.modify(SCTLR_EL1::M::Enable);
+	barrier::isb(barrier::SY);
 
 	info!("Successfully set up paging.");