feat(pvtime): Add PVTime device support for ARM

- Added PVTime struct and logic to allocate and register
  per-vCPU steal_time regions
- Integrated PVTime device initialization during microVM boot
- Fixed kvm_set_device_attr failure by aligning each
  steal_time region to 64 bytes
- Added snapshot/restore support via Persist interface
- Refactored PVTime into its own module

Signed-off-by: Dakshin Devanand <[email protected]>

Author: DakshinD

src/vmm/src/arch/aarch64/mod.rs

@@ -9,6 +9,8 @@ pub mod gic;
 pub mod kvm;
 /// Layout for this aarch64 system.
 pub mod layout;
+/// Module for the paravirtualized time device
+pub mod pvtime;
 /// Logic for configuring aarch64 registers.
 pub mod regs;
 /// Architecture specific vCPU code

src/vmm/src/arch/aarch64/pvtime.rs

@@ -0,0 +1,171 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+use displaydoc::Display;
+use kvm_bindings::{KVM_ARM_VCPU_PVTIME_CTRL, KVM_ARM_VCPU_PVTIME_IPA};
+use kvm_ioctls::VcpuFd;
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+use vm_memory::GuestAddress;
+
+use crate::Vcpu;
+use crate::device_manager::resources::ResourceAllocator;
+use crate::snapshot::Persist;
+
+/// 64 bytes due to alignment requirement in 3.1 of https://www.kernel.org/doc/html/v5.8/virt/kvm/devices/vcpu.html#attribute-kvm-arm-vcpu-pvtime-ipa
+pub const STEALTIME_STRUCT_MEM_SIZE: u64 = 64;
+
+/// Represent PVTime device for ARM
+#[derive(Debug)]
+pub struct PVTime {
+    /// Number of vCPUs
+    vcpu_count: u8,
+    /// The base IPA of the shared memory region
+    base_ipa: GuestAddress,
+}
+
+/// Errors associated with PVTime operations
+#[derive(Debug, Error, Display, PartialEq, Eq)]
+pub enum PVTimeError {
+    /// Failed to allocate memory region: {0}
+    AllocationFailed(vm_allocator::Error),
+    /// Invalid VCPU ID: {0}
+    InvalidVcpuIndex(u8),
+    /// Error while setting or getting device attributes for vCPU: {0}
+    DeviceAttribute(kvm_ioctls::Error),
+}
+
+impl PVTime {
+    /// Helper function to get the IPA of the steal_time region for a given vCPU
+    fn get_steal_time_region_addr(&self, vcpu_index: u8) -> Result<GuestAddress, PVTimeError> {
+        if vcpu_index >= self.vcpu_count {
+            return Err(PVTimeError::InvalidVcpuIndex(vcpu_index));
+        }
+        Ok(GuestAddress(
+            self.base_ipa.0 + (vcpu_index as u64 * STEALTIME_STRUCT_MEM_SIZE),
+        ))
+    }
+
+    /// Create a new PVTime device given a base addr
+    /// - Assumes total shared memory region from base addr is already allocated
+    fn from_base(base_ipa: GuestAddress, vcpu_count: u8) -> Self {
+        PVTime {
+            vcpu_count,
+            base_ipa,
+        }
+    }
+
+    /// Creates a new PVTime device by allocating new system memory for all vCPUs
+    pub fn new(
+        resource_allocator: &mut ResourceAllocator,
+        vcpu_count: u8,
+    ) -> Result<Self, PVTimeError> {
+        // This returns the IPA of the start of our shared memory region for all vCPUs.
+        let base_ipa: GuestAddress = GuestAddress(
+            resource_allocator
+                .allocate_system_memory(
+                    STEALTIME_STRUCT_MEM_SIZE * vcpu_count as u64,
+                    STEALTIME_STRUCT_MEM_SIZE,
+                    vm_allocator::AllocPolicy::LastMatch,
+                )
+                .map_err(PVTimeError::AllocationFailed)?,
+        );
+        Ok(Self::from_base(base_ipa, vcpu_count))
+    }
+
+    /// Check if PVTime is supported on vcpu
+    pub fn is_supported(vcpu_fd: &VcpuFd) -> bool {
+        // Check if pvtime is enabled
+        let pvtime_device_attr = kvm_bindings::kvm_device_attr {
+            group: kvm_bindings::KVM_ARM_VCPU_PVTIME_CTRL,
+            attr: kvm_bindings::KVM_ARM_VCPU_PVTIME_IPA as u64,
+            addr: 0,
+            flags: 0,
+        };
+
+        // Use kvm_has_device_attr to check if PVTime is supported
+        vcpu_fd.has_device_attr(&pvtime_device_attr).is_ok()
+    }
+
+    /// Register a vCPU with its pre-allocated steal time region
+    fn register_vcpu(
+        &self,
+        vcpu_index: u8,
+        vcpu_fd: &kvm_ioctls::VcpuFd,
+    ) -> Result<(), PVTimeError> {
+        // Get IPA of the steal_time region for this vCPU
+        let ipa = self.get_steal_time_region_addr(vcpu_index)?;
+
+        // Use KVM syscall (kvm_set_device_attr) to register the vCPU with the steal_time region
+        let vcpu_device_attr = kvm_bindings::kvm_device_attr {
+            group: KVM_ARM_VCPU_PVTIME_CTRL,
+            attr: KVM_ARM_VCPU_PVTIME_IPA as u64,
+            addr: &ipa.0 as *const u64 as u64, // userspace address of attr data
+            flags: 0,
+        };
+
+        vcpu_fd
+            .set_device_attr(&vcpu_device_attr)
+            .map_err(PVTimeError::DeviceAttribute)?;
+
+        Ok(())
+    }
+
+    /// Register all vCPUs with their pre-allocated steal time regions
+    pub fn register_all_vcpus(&self, vcpus: &mut [Vcpu]) -> Result<(), PVTimeError> {
+        // Register the vcpu with the pvtime device to map its steal time region
+        for (i, vcpu) in vcpus.iter().enumerate() {
+            #[allow(clippy::cast_possible_truncation)]
+            // We know vcpu_count is u8 according to VcpuConfig
+            self.register_vcpu(i as u8, &vcpu.kvm_vcpu.fd)?;
+        }
+        Ok(())
+    }
+}
+
+/// Logic to save/restore the state of a PVTime device
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub struct PVTimeState {
+    /// base IPA of the total shared memory region
+    pub base_ipa: u64,
+}
+
+/// Arguments to restore a PVTime device from PVTimeState
+#[derive(Debug)]
+pub struct PVTimeConstructorArgs<'a> {
+    /// For steal_time shared memory region
+    pub resource_allocator: &'a mut ResourceAllocator,
+    /// Number of vCPUs (should be consistent with pre-snapshot state)
+    pub vcpu_count: u8,
+}
+
+impl<'a> Persist<'a> for PVTime {
+    type State = PVTimeState;
+    type ConstructorArgs = PVTimeConstructorArgs<'a>;
+    type Error = PVTimeError;
+
+    /// Save base IPA of PVTime device for persistence
+    fn save(&self) -> Self::State {
+        PVTimeState {
+            base_ipa: self.base_ipa.0,
+        }
+    }
+
+    /// Restore state of PVTime device from given base IPA
+    fn restore(
+        constructor_args: Self::ConstructorArgs,
+        state: &Self::State,
+    ) -> std::result::Result<Self, PVTimeError> {
+        constructor_args
+            .resource_allocator
+            .allocate_system_memory(
+                STEALTIME_STRUCT_MEM_SIZE * constructor_args.vcpu_count as u64,
+                STEALTIME_STRUCT_MEM_SIZE,
+                vm_allocator::AllocPolicy::ExactMatch(state.base_ipa),
+            )
+            .map_err(PVTimeError::AllocationFailed)?;
+        Ok(Self::from_base(
+            GuestAddress(state.base_ipa),
+            constructor_args.vcpu_count,
+        ))
+    }
+}

src/vmm/src/builder.rs

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 //! Enables pre-boot setup, instantiation and booting of a Firecracker VMM.
-
 use std::fmt::Debug;
 use std::io;
 #[cfg(feature = "gdb")]
@@ -12,13 +11,16 @@ use std::sync::{Arc, Mutex};
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
+use log::warn;
 use userfaultfd::Uffd;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
 use vm_superio::Rtc;
 use vm_superio::Serial;
 use vmm_sys_util::eventfd::EventFd;
 
+#[cfg(target_arch = "aarch64")]
+use crate::arch::aarch64::pvtime::{PVTime, PVTimeConstructorArgs, PVTimeError};
 use crate::arch::{ConfigurationError, configure_system_for_boot, load_kernel};
 #[cfg(target_arch = "aarch64")]
 use crate::construct_kvm_mpidrs;
@@ -82,6 +84,8 @@ pub enum StartMicrovmError {
     CreateLegacyDevice(device_manager::legacy::LegacyDeviceError),
     /// Error creating VMGenID device: {0}
     CreateVMGenID(VmGenIdError),
+    /// Error creating PVTime device: {0}
+    CreatePVTime(PVTimeError),
     /// Invalid Memory Configuration: {0}
     GuestMemory(crate::vstate::memory::MemoryError),
     /// Error with initrd initialization: {0}.
@@ -189,6 +193,8 @@ fn create_vmm_and_vcpus(
         #[cfg(target_arch = "x86_64")]
         pio_device_manager,
         acpi_device_manager,
+        #[cfg(target_arch = "aarch64")]
+        pv_time: None,
     };
 
     Ok((vmm, vcpus))
@@ -289,6 +295,17 @@ pub fn build_microvm_for_boot(
 
     attach_vmgenid_device(&mut vmm)?;
 
+    // Attempt to setup PVTime, continue if not supported
+    #[cfg(target_arch = "aarch64")]
+    {
+        vmm.pv_time = if PVTime::is_supported(&vcpus[0].kvm_vcpu.fd) {
+            Some(setup_pv_time(&mut vmm, vcpus.as_mut())?)
+        } else {
+            warn!("PVTime is not supported by KVM. Steal time will not be reported to the guest.");
+            None
+        };
+    }
+
     configure_system_for_boot(
         &mut vmm,
         vcpus.as_mut(),
@@ -403,6 +420,8 @@ pub enum BuildMicrovmFromSnapshotError {
     ACPIDeviManager(#[from] ACPIDeviceManagerRestoreError),
     /// VMGenID update failed: {0}
     VMGenIDUpdate(std::io::Error),
+    /// Failed to restore PVTime device: {0}
+    RestorePVTime(#[from] PVTimeError),
 }
 
 /// Builds and starts a microVM based on the provided MicrovmState.
@@ -453,6 +472,29 @@ pub fn build_microvm_from_snapshot(
             .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
     }
 
+    // Restore the PVTime device
+    #[cfg(target_arch = "aarch64")]
+    {
+        let pvtime_state = microvm_state.pvtime_state;
+        if let Some(pvtime_state) = pvtime_state {
+            #[allow(clippy::cast_possible_truncation)]
+            // We know vcpu_count is u8 according to VcpuConfig
+            let pvtime_ctor_args = PVTimeConstructorArgs {
+                resource_allocator: &mut vmm.resource_allocator,
+                vcpu_count: vcpus.len() as u8,
+            };
+            vmm.pv_time = Some(
+                PVTime::restore(pvtime_ctor_args, &pvtime_state)
+                    .map_err(BuildMicrovmFromSnapshotError::RestorePVTime)?,
+            );
+            vmm.pv_time
+                .as_ref()
+                .unwrap()
+                .register_all_vcpus(&mut vcpus) // We can safely unwrap here
+                .map_err(StartMicrovmError::CreatePVTime)?;
+        }
+    }
+
     #[cfg(target_arch = "aarch64")]
     {
         let mpidrs = construct_kvm_mpidrs(&microvm_state.vcpu_states);
@@ -548,6 +590,24 @@ pub fn setup_serial_device(
     Ok(serial)
 }
 
+/// Sets up the pvtime device.
+#[cfg(target_arch = "aarch64")]
+fn setup_pv_time(vmm: &mut Vmm, vcpus: &mut [Vcpu]) -> Result<PVTime, StartMicrovmError> {
+    use crate::arch::aarch64::pvtime::PVTime;
+
+    // Create the pvtime device
+    #[allow(clippy::cast_possible_truncation)] // We know vcpu_count is u8 according to VcpuConfig
+    let pv_time = PVTime::new(&mut vmm.resource_allocator, vcpus.len() as u8)
+        .map_err(StartMicrovmError::CreatePVTime)?;
+
+    // Register all vcpus with pvtime device
+    pv_time
+        .register_all_vcpus(vcpus)
+        .map_err(StartMicrovmError::CreatePVTime)?;
+
+    Ok(pv_time)
+}
+
 #[cfg(target_arch = "aarch64")]
 fn attach_legacy_devices_aarch64(
     event_manager: &mut EventManager,
@@ -851,6 +911,8 @@ pub(crate) mod tests {
             #[cfg(target_arch = "x86_64")]
             pio_device_manager,
             acpi_device_manager,
+            #[cfg(target_arch = "aarch64")]
+            pv_time: None,
         }
     }
 

src/vmm/src/lib.rs

@@ -134,6 +134,8 @@ use vstate::kvm::Kvm;
 use vstate::vcpu::{self, StartThreadedError, VcpuSendEventError};
 
 use crate::arch::DeviceType;
+#[cfg(target_arch = "aarch64")]
+use crate::arch::aarch64::pvtime::PVTime;
 use crate::cpu_config::templates::CpuConfiguration;
 #[cfg(target_arch = "x86_64")]
 use crate::device_manager::legacy::PortIODeviceManager;
@@ -157,7 +159,6 @@ use crate::vstate::memory::{
 use crate::vstate::vcpu::VcpuState;
 pub use crate::vstate::vcpu::{Vcpu, VcpuConfig, VcpuEvent, VcpuHandle, VcpuResponse};
 pub use crate::vstate::vm::Vm;
-
 /// Shorthand type for the EventManager flavour used by Firecracker.
 pub type EventManager = BaseEventManager<Arc<Mutex<dyn MutEventSubscriber>>>;
 
@@ -324,6 +325,8 @@ pub struct Vmm {
     #[cfg(target_arch = "x86_64")]
     pio_device_manager: PortIODeviceManager,
     acpi_device_manager: ACPIDeviceManager,
+    #[cfg(target_arch = "aarch64")]
+    pv_time: Option<PVTime>,
 }
 
 impl Vmm {
@@ -523,6 +526,9 @@ impl Vmm {
 
         let memory_state = self.guest_memory.describe();
         let acpi_dev_state = self.acpi_device_manager.save();
+        #[cfg(target_arch = "aarch64")]
+        let pvtime_state: Option<arch::aarch64::pvtime::PVTimeState> =
+            self.pv_time.as_ref().map(|pvtime| pvtime.save());
 
         Ok(MicrovmState {
             vm_info: vm_info.clone(),
@@ -532,6 +538,8 @@ impl Vmm {
             vcpu_states,
             device_states,
             acpi_dev_state,
+            #[cfg(target_arch = "aarch64")]
+            pvtime_state,
         })
     }
 

src/vmm/src/persist.rs

@@ -16,6 +16,8 @@ use serde::{Deserialize, Serialize};
 use userfaultfd::{FeatureFlags, Uffd, UffdBuilder};
 use vmm_sys_util::sock_ctrl_msg::ScmSocket;
 
+#[cfg(target_arch = "aarch64")]
+use crate::arch::aarch64::pvtime::PVTimeState;
 #[cfg(target_arch = "aarch64")]
 use crate::arch::aarch64::vcpu::get_manufacturer_id_from_host;
 use crate::builder::{self, BuildMicrovmFromSnapshotError};
@@ -88,6 +90,9 @@ pub struct MicrovmState {
     pub device_states: DeviceStates,
     /// ACPI devices state.
     pub acpi_dev_state: ACPIDeviceManagerState,
+    /// PVTime device state (optional for platforms that support it).
+    #[cfg(target_arch = "aarch64")]
+    pub pvtime_state: Option<PVTimeState>,
 }
 
 /// This describes the mapping between Firecracker base virtual address and
@@ -765,6 +770,8 @@ mod tests {
             #[cfg(target_arch = "x86_64")]
             vm_state: vmm.vm.save_state().unwrap(),
             acpi_dev_state: vmm.acpi_device_manager.save(),
+            #[cfg(target_arch = "aarch64")]
+            pvtime_state: None,
         };
 
         let mut buf = vec![0; 10000];