vmbase: Support 16KiB MMIO_GUARD granule

Now that the MMIO_GUARD sharing/unsharing is decoupled from the page
tables, add support for 16KiB (and possibly beyond) by getting rid of
the hard-coded assumption that MMIO_GUARD happens with 4KiB granularity
by expecting potentially more than one page fault per MMIO_GUARD region
with MMIO lazy mapping.

Address the special case of the UART separately (see comment there).

Remove the now-obsolete hard-coded MMIO_GUARD_GRANULE_SIZE.

Bug: 336563593
Test: m libpvmfw libvmbase_example librialto
Change-Id: I5d8e169d9c2ad022208de3d33f346f5034d13393
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index 892dc9f..bb5ccef 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -15,13 +15,13 @@
 //! Rust entry point.
 
 use crate::{
-    bionic, console, heap,
-    hyp::{self, MMIO_GUARD_GRANULE_SIZE},
-    logger,
+    bionic, console, heap, hyp, logger,
+    memory::{page_4kb_of, SIZE_16KB, SIZE_4KB},
     power::{reboot, shutdown},
     rand,
 };
 use core::mem::size_of;
+use static_assertions::const_assert_eq;
 
 fn try_console_init() -> Result<(), hyp::Error> {
     console::init();
@@ -31,8 +31,19 @@
 
         // TODO(ptosi): Use MmioSharer::share() to properly track this MMIO_GUARD_MAP.
         //
-        // MmioSharer only supports MMIO_GUARD_GRANULE_SIZE so fail early here if needed.
-        assert_eq!(mmio_guard.granule()?, MMIO_GUARD_GRANULE_SIZE);
+        // The following call shares the UART but also anything else present in 0..granule.
+        //
+        // For 4KiB, that's only the UARTs. For 16KiB, it also covers the RTC and watchdog but, as
+        // neither is used by vmbase clients (and as both are outside of the UART page), they
+        // will never have valid stage-1 mappings to those devices. As a result, this
+        // MMIO_GUARD_MAP isn't affected by the granule size in any visible way. Larger granule
+        // sizes will need to be checked separately, if needed.
+        assert!({
+            let granule = mmio_guard.granule()?;
+            granule == SIZE_4KB || granule == SIZE_16KB
+        });
+        // Validate the assumption above by ensuring that the UART is not moved to another page:
+        const_assert_eq!(page_4kb_of(console::BASE_ADDRESS), 0);
         mmio_guard.map(console::BASE_ADDRESS)?;
     }
 
diff --git a/vmbase/src/hyp.rs b/vmbase/src/hyp.rs
index 1461f80..1cc2ca7 100644
--- a/vmbase/src/hyp.rs
+++ b/vmbase/src/hyp.rs
@@ -20,5 +20,4 @@
 pub use error::{Error, Result};
 pub use hypervisor::{
     get_device_assigner, get_mem_sharer, get_mmio_guard, DeviceAssigningHypervisor, KvmError,
-    MMIO_GUARD_GRANULE_SIZE,
 };
diff --git a/vmbase/src/hyp/hypervisor.rs b/vmbase/src/hyp/hypervisor.rs
index dab15ec..1b45f38 100644
--- a/vmbase/src/hyp/hypervisor.rs
+++ b/vmbase/src/hyp/hypervisor.rs
@@ -22,9 +22,7 @@
 use super::{Error, Result};
 use alloc::boxed::Box;
 use common::Hypervisor;
-pub use common::{
-    DeviceAssigningHypervisor, MemSharingHypervisor, MmioGuardedHypervisor, MMIO_GUARD_GRANULE_SIZE,
-};
+pub use common::{DeviceAssigningHypervisor, MemSharingHypervisor, MmioGuardedHypervisor};
 pub use geniezone::GeniezoneError;
 use geniezone::GeniezoneHypervisor;
 use gunyah::GunyahHypervisor;
diff --git a/vmbase/src/hyp/hypervisor/common.rs b/vmbase/src/hyp/hypervisor/common.rs
index 1ed0696..de0fe12 100644
--- a/vmbase/src/hyp/hypervisor/common.rs
+++ b/vmbase/src/hyp/hypervisor/common.rs
@@ -14,10 +14,7 @@
 
 //! This module regroups some common traits shared by all the hypervisors.
 
-use crate::{hyp::Result, memory::SIZE_4KB};
-
-/// Expected MMIO guard granule size, validated during MMIO guard initialization.
-pub const MMIO_GUARD_GRANULE_SIZE: usize = SIZE_4KB;
+use crate::hyp::Result;
 
 /// Trait for the hypervisor.
 pub trait Hypervisor {
diff --git a/vmbase/src/memory.rs b/vmbase/src/memory.rs
index 2f72fc4..299d50f 100644
--- a/vmbase/src/memory.rs
+++ b/vmbase/src/memory.rs
@@ -26,8 +26,8 @@
     handle_permission_fault, handle_translation_fault, MemoryRange, MemoryTracker, MEMORY,
 };
 pub use util::{
-    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_2MB,
-    SIZE_4KB, SIZE_4MB, SIZE_64KB,
+    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_16KB,
+    SIZE_2MB, SIZE_4KB, SIZE_4MB, SIZE_64KB,
 };
 
 pub(crate) use shared::{alloc_shared, dealloc_shared};
diff --git a/vmbase/src/memory/shared.rs b/vmbase/src/memory/shared.rs
index 457e6f2..5a25d9f 100644
--- a/vmbase/src/memory/shared.rs
+++ b/vmbase/src/memory/shared.rs
@@ -21,7 +21,7 @@
 use crate::console;
 use crate::dsb;
 use crate::exceptions::HandleExceptionError;
-use crate::hyp::{self, get_mem_sharer, get_mmio_guard, MMIO_GUARD_GRANULE_SIZE};
+use crate::hyp::{self, get_mem_sharer, get_mmio_guard};
 use crate::util::unchecked_align_down;
 use crate::util::RangeExt as _;
 use aarch64_paging::paging::{
@@ -42,7 +42,6 @@
 use log::{debug, error, trace};
 use once_cell::race::OnceBox;
 use spin::mutex::SpinMutex;
-use static_assertions::const_assert_eq;
 use tinyvec::ArrayVec;
 
 /// A global static variable representing the system memory tracker, protected by a spin mutex.
@@ -397,12 +396,11 @@
     }
 
     fn get_granule() -> Result<usize> {
-        const_assert_eq!(MMIO_GUARD_GRANULE_SIZE, PAGE_SIZE); // For good measure.
         let Some(mmio_guard) = get_mmio_guard() else {
             return Ok(PAGE_SIZE);
         };
         match mmio_guard.granule()? {
-            MMIO_GUARD_GRANULE_SIZE => Ok(MMIO_GUARD_GRANULE_SIZE),
+            granule if granule % PAGE_SIZE == 0 => Ok(granule), // For good measure.
             granule => Err(MemoryTrackerError::UnsupportedMmioGuardGranule(granule)),
         }
     }
diff --git a/vmbase/src/memory/util.rs b/vmbase/src/memory/util.rs
index 2b75414..e9f867f 100644
--- a/vmbase/src/memory/util.rs
+++ b/vmbase/src/memory/util.rs
@@ -22,6 +22,8 @@
 
 /// The size of a 4KB memory in bytes.
 pub const SIZE_4KB: usize = 4 << 10;
+/// The size of a 16KB memory in bytes.
+pub const SIZE_16KB: usize = 16 << 10;
 /// The size of a 64KB memory in bytes.
 pub const SIZE_64KB: usize = 64 << 10;
 /// The size of a 128KB memory in bytes.