fix(aarch64): re-apply CLIDR_EL1 after PSCI CPU_ON

CHANGELOG.md

@@ -71,6 +71,11 @@ and this project adheres to
 
 ### Fixed
 
+- Fixed aarch64 SMP guests losing the DT-consistent `CLIDR_EL1` override
+  on secondary vCPUs after PSCI `CPU_ON`. KVM resets those vCPUs and
+  restores a fabricated `CLIDR_EL1`, which made guest cache topology
+  asymmetric and disabled Linux load balancing. The same override is now
+  written again after vCPU reset and when a secondary is powered on.
 - [#6100](https://github.com/firecracker-microvm/firecracker/pull/6100): Fixed
   the vsock device permanently suppressing RX (host-to-guest) delivery after a
   bare pause/resume cycle (`PATCH /vm` with `Paused` then `Resumed`, without a

src/vmm/src/arch/aarch64/cache_info.rs

@@ -413,6 +413,23 @@ pub(crate) fn merge_clidr(current: u64, sysfs: u64) -> u64 {
     (current & !REPLACE_MASK) | (sysfs & REPLACE_MASK)
 }
 
+/// Compute the DT-consistent CLIDR_EL1 value that should be written to each
+/// vCPU, or `None` when the override must be skipped.
+///
+/// Skip when sysfs reports no L1 caches (writing 0 would be worse than KVM's
+/// fabricated value) or when the merged value already matches `current`.
+pub(crate) fn clidr_override_from_current(current: u64) -> Result<Option<u64>, CacheInfoError> {
+    let mut l1_caches = Vec::new();
+    let mut non_l1_caches = Vec::new();
+    read_cache_config(&mut l1_caches, &mut non_l1_caches)?;
+    if l1_caches.is_empty() {
+        warn!("No L1 caches found in sysfs, skipping CLIDR override");
+        return Ok(None);
+    }
+    let new_clidr = merge_clidr(current, build_clidr_from_caches(&l1_caches, &non_l1_caches));
+    Ok((new_clidr != current).then_some(new_clidr))
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
@@ -772,4 +789,19 @@ mod tests {
         let sysfs = 0x0000_0000_0300_0123_u64;
         assert_eq!(merge_clidr(current, sysfs), current);
     }
+
+    #[test]
+    fn test_clidr_override_from_current() {
+        // Mock store: L1 Data + L1 Instruction + L2 Unified (ctype1=3, ctype2=4, LoC=2).
+        let current_fabricated: u64 = 4; // unified L1 only
+        let override_val = clidr_override_from_current(current_fabricated)
+            .unwrap()
+            .expect("fabricated CLIDR should be overridden from mock sysfs");
+        assert_eq!(override_val & 0x7, 3, "L1 should be Separate");
+        assert_eq!((override_val >> 3) & 0x7, 4, "L2 should be Unified");
+        assert_eq!((override_val >> 24) & 0x7, 2, "LoC should be 2");
+
+        // Already-matching current is a no-op.
+        assert_eq!(clidr_override_from_current(override_val).unwrap(), None);
+    }
 }

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

@@ -28,7 +28,6 @@ use crate::arch::{BootProtocol, EntryPoint, arch_memory_regions_with_gap};
 use crate::cpu_config::aarch64::{CpuConfiguration, CpuConfigurationError};
 use crate::cpu_config::templates::CustomCpuTemplate;
 use crate::initrd::InitrdConfig;
-use zerocopy::IntoBytes;
 
 use crate::logger::warn;
 use crate::utils::{u64_to_usize, usize_to_u64};
@@ -124,7 +123,8 @@ pub fn configure_system_for_boot(
     }
 
     // Override CLIDR_EL1 ctype/LoC fields on each vCPU to match the host's
-    // real cache topology. See `override_clidr` for details.
+    // real cache topology. See `override_clidr` for details. The value is
+    // stored on each vCPU so it can be written again after PSCI CPU_ON.
     override_clidr(vcpus)?;
 
     let vcpu_mpidr = vcpus
@@ -166,49 +166,29 @@ pub fn configure_system_for_boot(
 /// and LoC fields with values derived from sysfs, and preserve all other fields
 /// (LoUU, LoUIS, ICB, Ttype). This is safe on pre-6.3 kernels where CLIDR
 /// already matches sysfs — the write is skipped as a no-op.
-fn override_clidr(vcpus: &[Vcpu]) -> Result<(), ConfigurationError> {
-    let mut l1_caches = Vec::new();
-    let mut non_l1_caches = Vec::new();
-    cache_info::read_cache_config(&mut l1_caches, &mut non_l1_caches)?;
-
-    // If sysfs reports no L1 caches, we cannot build a meaningful CLIDR.
-    // Writing an all-zero CLIDR would tell the guest there are no caches,
-    // which is worse than whatever KVM fabricated. Leave it alone.
-    if l1_caches.is_empty() {
-        warn!("No L1 caches found in sysfs, skipping CLIDR override");
-        return Ok(());
-    }
-
-    let sysfs_clidr = cache_info::build_clidr_from_caches(&l1_caches, &non_l1_caches);
-
-    let mut cur_clidr: u64 = 0;
+fn override_clidr(vcpus: &mut [Vcpu]) -> Result<(), ConfigurationError> {
     // Reading/writing CLIDR_EL1 via KVM_SET_ONE_REG may not be supported on
     // older kernels (pre-6.3). In that case KVM passes through the real host
     // CLIDR and the override is unnecessary, so we warn and continue.
-    if let Err(e) = vcpus[0]
-        .kvm_vcpu
-        .fd
-        .get_one_reg(regs::CLIDR_EL1, cur_clidr.as_mut_bytes())
-    {
-        warn!("Failed to read CLIDR_EL1, skipping override: {e}");
+    let cur_clidr = match vcpus[0].kvm_vcpu.get_clidr() {
+        Ok(value) => value,
+        Err(e) => {
+            warn!("Failed to read CLIDR_EL1, skipping override: {e}");
+            return Ok(());
+        }
+    };
+
+    let Some(new_clidr) = cache_info::clidr_override_from_current(cur_clidr)? else {
         return Ok(());
-    }
+    };
 
-    let new_clidr = cache_info::merge_clidr(cur_clidr, sysfs_clidr);
-
-    if new_clidr != cur_clidr {
-        for vcpu in vcpus.iter() {
-            if let Err(e) = vcpu
-                .kvm_vcpu
-                .fd
-                .set_one_reg(regs::CLIDR_EL1, new_clidr.as_bytes())
-            {
-                warn!(
-                    "Failed to set CLIDR_EL1 to {:#x} on vCPU {}, skipping override: {e}",
-                    new_clidr, vcpu.kvm_vcpu.index
-                );
-                return Ok(());
-            }
+    for vcpu in vcpus.iter_mut() {
+        if let Err(e) = vcpu.kvm_vcpu.apply_clidr_override(new_clidr) {
+            warn!(
+                "Failed to set CLIDR_EL1 to {:#x} on vCPU {}, skipping override: {e}",
+                new_clidr, vcpu.kvm_vcpu.index
+            );
+            return Ok(());
         }
     }
 

src/vmm/src/arch/aarch64/vcpu.rs

@@ -117,6 +117,12 @@ pub struct KvmVcpu {
     kvi: kvm_vcpu_init,
     /// IPA of steal_time region
     pub pvtime_ipa: Option<GuestAddress>,
+    /// DT-consistent CLIDR_EL1 written at boot. Re-applied after vCPU reset
+    /// (KVM_ARM_VCPU_INIT and in-kernel PSCI CPU_ON).
+    clidr_override: Option<u64>,
+    /// Whether [`Self::prepare_clidr_for_run`] has already restored
+    /// [`Self::clidr_override`] after this secondary was powered on.
+    clidr_reapplied_after_power_on: bool,
 }
 
 /// Vcpu peripherals
@@ -151,9 +157,65 @@ impl KvmVcpu {
             peripherals: Default::default(),
             kvi,
             pvtime_ipa: None,
+            clidr_override: None,
+            clidr_reapplied_after_power_on: false,
         })
     }
 
+    /// Read CLIDR_EL1.
+    pub fn get_clidr(&self) -> Result<u64, VcpuArchError> {
+        let mut clidr = [0_u8; 8];
+        self.fd
+            .get_one_reg(CLIDR_EL1, &mut clidr)
+            .map_err(|err| VcpuArchError::GetOneReg(CLIDR_EL1, err))?;
+        Ok(u64::from_le_bytes(clidr))
+    }
+
+    /// Write CLIDR_EL1 without recording it as the persistent override.
+    fn set_clidr(&self, value: u64) -> Result<(), VcpuArchError> {
+        self.fd
+            .set_one_reg(CLIDR_EL1, &value.to_le_bytes())
+            .map_err(|err| VcpuArchError::SetOneReg(CLIDR_EL1, format!("{value:#x}"), err))
+    }
+
+    /// Store and write the DT-consistent CLIDR_EL1 override for this vCPU.
+    pub fn apply_clidr_override(&mut self, value: u64) -> Result<(), VcpuArchError> {
+        self.clidr_override = Some(value);
+        self.set_clidr(value)
+    }
+
+    /// Write the stored CLIDR_EL1 override again, if one was configured.
+    pub fn reapply_clidr_override(&self) -> Result<(), VcpuArchError> {
+        if let Some(value) = self.clidr_override {
+            self.set_clidr(value)?;
+        }
+        Ok(())
+    }
+
+    /// Restore the CLIDR_EL1 override after in-kernel PSCI CPU_ON.
+    ///
+    /// Secondary vCPUs are created powered-off. PSCI CPU_ON makes KVM reset
+    /// the target and run `reset_clidr()`, which overwrites userspace's
+    /// SET_ONE_REG. Once the vCPU is no longer STOPPED, this writes the
+    /// override. KVM folds a pending `KVM_REQ_VCPU_RESET` into architected
+    /// state before SET_ONE_REG, so the write lands after reset.
+    ///
+    /// Returns `false` while the vCPU is still powered off so the caller can
+    /// service vCPU events instead of blocking in a poll loop.
+    pub fn prepare_clidr_for_run(&mut self) -> Result<bool, VcpuArchError> {
+        if self.clidr_override.is_none() || self.index == 0 || self.clidr_reapplied_after_power_on {
+            return Ok(true);
+        }
+
+        if self.get_mpstate()?.mp_state == KVM_MP_STATE_STOPPED {
+            return Ok(false);
+        }
+
+        self.reapply_clidr_override()?;
+        self.clidr_reapplied_after_power_on = true;
+        Ok(true)
+    }
+
     /// Read the MPIDR - Multiprocessor Affinity Register.
     pub fn get_mpidr(&self) -> Result<u64, VcpuArchError> {
         // MPIDR register is 64 bit wide on aarch64
@@ -312,6 +374,11 @@ impl KvmVcpu {
         }
 
         self.fd.vcpu_init(&self.kvi).map_err(KvmVcpuError::Init)?;
+        // KVM_ARM_VCPU_INIT resets sysregs, including CLIDR_EL1. Restore the
+        // DT-consistent override so a secondary that is reset (the same
+        // sequence as PSCI CPU_ON) does not keep KVM's fabricated value.
+        self.reapply_clidr_override()
+            .map_err(KvmVcpuError::ConfigureRegisters)?;
         Ok(())
     }
 
@@ -550,7 +617,7 @@ mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
     use std::os::unix::io::AsRawFd;
 
-    use kvm_bindings::{KVM_ARM_VCPU_PSCI_0_2, KVM_REG_SIZE_U64};
+    use kvm_bindings::{KVM_ARM_VCPU_PSCI_0_2, KVM_MP_STATE_RUNNABLE, KVM_REG_SIZE_U64, kvm_mp_state};
     use vm_memory::GuestAddress;
 
     use super::*;
@@ -787,6 +854,82 @@ mod tests {
         vcpu2.init(&[]).unwrap();
     }
 
+    /// Write a sentinel CLIDR_EL1 that is distinct from KVM's reset value.
+    /// Returns `None` when the register is not readable or not writable
+    /// (pre-6.3 kernels pass through the host value and reject the override).
+    fn try_install_clidr_sentinel(vcpu: &mut KvmVcpu) -> Option<u64> {
+        let original = vcpu.get_clidr().ok()?;
+        let override_val = original ^ 0x7;
+        vcpu.apply_clidr_override(override_val).ok()?;
+        (vcpu.get_clidr().ok()? == override_val).then_some(override_val)
+    }
+
+    #[test]
+    fn test_clidr_override_survives_secondary_vcpu_reset() {
+        let vm = setup_vm_with_memory(0x1000);
+        let mut vcpu0 = KvmVcpu::new(0, &vm).unwrap();
+        let mut vcpu1 = KvmVcpu::new(1, &vm).unwrap();
+        vcpu0.init(&[]).unwrap();
+        vcpu1.init(&[]).unwrap();
+
+        let Some(override_val) = try_install_clidr_sentinel(&mut vcpu1) else {
+            return;
+        };
+        vcpu0.apply_clidr_override(override_val).unwrap();
+
+        // KVM_ARM_VCPU_INIT resets sysregs the same way PSCI CPU_ON does
+        // (`kvm_reset_vcpu` → `reset_clidr()`). The Firecracker path that
+        // wraps that ioctl must write the override again; otherwise only
+        // vCPU0 (which is not reset) keeps the DT-consistent value.
+        vcpu1.init_vcpu().unwrap();
+
+        assert_eq!(
+            vcpu1.get_clidr().unwrap(),
+            override_val,
+            "secondary vCPU lost CLIDR_EL1 override after reset"
+        );
+        assert_eq!(
+            vcpu0.get_clidr().unwrap(),
+            override_val,
+            "vCPU0 CLIDR_EL1 override must be unchanged"
+        );
+    }
+
+    #[test]
+    fn test_clidr_override_reapplied_after_secondary_power_on() {
+        let vm = setup_vm_with_memory(0x1000);
+        let mut vcpu1 = KvmVcpu::new(1, &vm).unwrap();
+        vcpu1.init(&[]).unwrap();
+
+        let Some(override_val) = try_install_clidr_sentinel(&mut vcpu1) else {
+            return;
+        };
+
+        // Raw KVM reset, like in-kernel PSCI CPU_ON: clobbers userspace CLIDR
+        // without going through `init_vcpu`.
+        vcpu1.fd.vcpu_init(&vcpu1.kvi).unwrap();
+
+        // Still powered off: do not enter KVM_RUN, so vCPU events stay live.
+        assert!(
+            !vcpu1.prepare_clidr_for_run().unwrap(),
+            "powered-off secondary must not be treated as ready for KVM_RUN"
+        );
+
+        // CPU_ON makes the target runnable. SET_ONE_REG then folds any
+        // pending KVM_REQ_VCPU_RESET and writes the override after it.
+        let mp = kvm_mp_state {
+            mp_state: KVM_MP_STATE_RUNNABLE,
+        };
+        vcpu1.set_mpstate(mp).unwrap();
+
+        assert!(vcpu1.prepare_clidr_for_run().unwrap());
+        assert_eq!(
+            vcpu1.get_clidr().unwrap(),
+            override_val,
+            "CLIDR_EL1 override must be restored after secondary power-on"
+        );
+    }
+
     #[test]
     fn test_get_valid_regs() {
         // Test `get_regs()` with valid register IDs.

src/vmm/src/vstate/vcpu.rs

@@ -411,6 +411,21 @@ impl Vcpu {
             return Ok(VcpuEmulation::Interrupted);
         }
 
+        // Secondary aarch64 vCPUs are powered off until PSCI CPU_ON. That
+        // in-kernel path resets CLIDR_EL1; restore the boot override after
+        // the vCPU becomes runnable and before the first guest entry.
+        #[cfg(target_arch = "aarch64")]
+        match self.kvm_vcpu.prepare_clidr_for_run() {
+            Ok(true) => {}
+            Ok(false) => {
+                thread::sleep(Duration::from_millis(1));
+                return Ok(VcpuEmulation::Interrupted);
+            }
+            Err(err) => {
+                warn!("Failed to restore CLIDR_EL1 after secondary power-on: {err}");
+            }
+        }
+
         match self.kvm_vcpu.fd.run() {
             Err(ref err) if err.errno() == libc::EINTR => {
                 self.kvm_vcpu.fd.set_kvm_immediate_exit(0);

tests/integration_tests/functional/test_topology.py

@@ -169,7 +169,27 @@ def parse_cache_info(info: str):
     return guest_cache_info, host_cache_info
 
 
+def _aarch64_assert_symmetric_cache_leaves(test_microvm, no_cpus):
+    """Every vCPU must expose the same cache index set.
+
+    A CLIDR_EL1 override that sticks only on vCPU0 leaves secondaries with
+    KVM's fabricated topology. Linux then fails to build scheduler domains.
+    """
+    cmd = (
+        f"for i in $(seq 0 {no_cpus - 1}); do "
+        f'printf "cpu%s %s\\n" "$i" '
+        f'"$(ls /sys/devices/system/cpu/cpu$i/cache 2>/dev/null | grep -c "^index" || true)"; '
+        f"done"
+    )
+    _, stdout, stderr = test_microvm.ssh.run(cmd)
+    assert stderr == ""
+    counts = [int(line.split()[1]) for line in stdout.splitlines() if line.strip()]
+    assert counts, "no cache index counts from guest"
+    assert len(set(counts)) == 1, f"asymmetric guest cache leaves per vCPU: {counts}"
+
+
 def _check_cache_topology_arm(test_microvm, no_cpus, kernel_version_tpl):
+    _aarch64_assert_symmetric_cache_leaves(test_microvm, no_cpus)
     guest_cache_info, host_cache_info = _aarch64_parse_cache_info(test_microvm, no_cpus)
 
     # Starting from 6.3 kernel cache representation for aarch64 platform has changed.