Batch register restore on MSHV and WHP

src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs

@@ -232,6 +232,36 @@ impl VirtualMachine for FaultInjectingVirtualMachine {
         self.inner().set_xcr0(value)
     }
 
+    #[cfg(target_arch = "x86_64")]
+    fn can_batch_registers(&self) -> bool {
+        self.inner().can_batch_registers()
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    fn set_batched_registers(
+        &mut self,
+        regs: &CommonRegisters,
+        debug_regs: &CommonDebugRegs,
+        sregs: &CommonSpecialRegisters,
+        xcr0: u64,
+        msrs: &[MsrEntry],
+    ) -> std::result::Result<(), RegisterError> {
+        if self.should_fail(VmOperation::SetRegs) {
+            return Err(RegisterError::SetRegs(Self::injected_error()));
+        }
+        if self.should_fail(VmOperation::SetDebugRegs) {
+            return Err(RegisterError::SetDebugRegs(Self::injected_error()));
+        }
+        if self.should_fail(VmOperation::SetSregs) {
+            return Err(RegisterError::SetSregs(Self::injected_error()));
+        }
+        if self.should_fail(VmOperation::SetMsrs) {
+            return Err(RegisterError::SetMsrs(Self::injected_error()));
+        }
+        self.inner_mut()
+            .set_batched_registers(regs, debug_regs, sregs, xcr0, msrs)
+    }
+
     #[cfg(target_arch = "aarch64")]
     fn can_reset_vcpu(&self) -> bool {
         self.inner().can_reset_vcpu()

src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs

@@ -239,22 +239,13 @@ impl HyperlightVm {
         Ok(self.vm.msrs(&self.msr_reset.persist_indices())?)
     }
 
-    /// Restores snapshot MSRs or the initialization baseline.
+    /// Restores snapshot MSRs and resets all other MSRs to their initialization values.
     pub(crate) fn restore_msrs(
         &mut self,
-        snap_msrs: Option<&Vec<MsrEntry>>,
+        snap_msrs: &[MsrEntry],
     ) -> std::result::Result<(), ResetVcpuError> {
-        match snap_msrs {
-            // No captured MSRs. Use this VM's baseline.
-            None => self.vm.set_msrs(self.msr_reset.baseline())?,
-            // Scrub the reset set to the destination baseline and write the
-            // snapshot's captured values on top. Validation rejects any
-            // captured index the destination cannot restore.
-            Some(msrs) => {
-                let entries = self.msr_reset.validate_snapshot(msrs)?;
-                self.vm.set_msrs(&entries)?;
-            }
-        }
+        let msrs = self.msr_reset.validate_snapshot(snap_msrs)?;
+        self.vm.set_msrs(&msrs)?;
         Ok(())
     }
 
@@ -320,41 +311,65 @@ impl HyperlightVm {
     /// - XSAVE (includes FPU/SSE state with proper FCW and MXCSR defaults)
     /// - XCR0
     /// - Special registers (restored from snapshot, with CR3 updated to new page table location)
+    /// - Model-specific registers (restored from snapshot, with omitted values reset to their
+    ///   initialization values)
     // TODO: check if other state needs to be reset
     pub(crate) fn reset_vcpu(
         &mut self,
         cr3: u64,
         sregs: &CommonSpecialRegisters,
+        snapshot_msrs: &[MsrEntry],
     ) -> std::result::Result<(), ResetVcpuError> {
-        self.vm.set_regs(&CommonRegisters {
+        let regs = CommonRegisters {
             rflags: 1 << 1, // Reserved bit always set
             ..Default::default()
-        })?;
-        self.vm.set_debug_regs(&CommonDebugRegs::default())?;
-        self.vm.reset_xsave()?;
-        self.vm.set_xcr0(XCR0_RESET)?;
+        };
+        let debug_regs = CommonDebugRegs::default();
+        let sregs = Self::sregs_with_cr3(cr3, sregs)?;
+        let msrs = self.msr_reset.validate_snapshot(snapshot_msrs)?;
 
-        self.apply_sregs(cr3, sregs)?;
+        self.pending_tlb_flush = true;
+        // Batch to avoid multiple hvcall overhead if supported
+        if self.vm.can_batch_registers() {
+            self.vm.reset_xsave()?;
+            self.vm
+                .set_batched_registers(&regs, &debug_regs, &sregs, XCR0_RESET, &msrs)?;
+        } else {
+            self.vm.set_regs(&regs)?;
+            self.vm.set_debug_regs(&debug_regs)?;
+            self.vm.reset_xsave()?;
+            self.vm.set_xcr0(XCR0_RESET)?;
+            self.vm.set_sregs(&sregs)?;
+            self.vm.set_msrs(&msrs)?;
+        }
 
         Ok(())
     }
 
-    /// Apply special registers and mark TLB for flush.
-    pub(crate) fn apply_sregs(
-        &mut self,
+    fn sregs_with_cr3(
         cr3: u64,
         sregs: &CommonSpecialRegisters,
-    ) -> std::result::Result<(), RegisterError> {
+    ) -> std::result::Result<CommonSpecialRegisters, RegisterError> {
         if sregs.apic_base & crate::hypervisor::regs::APIC_BASE_X2APIC_ENABLE != 0 {
             return Err(RegisterError::InvalidSnapshotApicBase {
                 value: sregs.apic_base,
             });
         }
 
-        // Restore the full special registers from snapshot, but update CR3
-        // to point to the new (relocated) page tables
         let mut sregs = *sregs;
         sregs.cr3 = cr3;
+        Ok(sregs)
+    }
+
+    /// Apply special registers and mark TLB for flush.
+    pub(crate) fn apply_sregs(
+        &mut self,
+        cr3: u64,
+        sregs: &CommonSpecialRegisters,
+    ) -> std::result::Result<(), RegisterError> {
+        // Restore the full special registers from snapshot, but update CR3
+        // to point to the new (relocated) page tables
+        let sregs = Self::sregs_with_cr3(cr3, sregs)?;
         self.pending_tlb_flush = true;
         self.vm.set_sregs(&sregs)?;
 
@@ -577,11 +592,7 @@ impl HyperlightVm {
     /// Tests use it to classify each resettable MSR.
     #[cfg(test)]
     pub(crate) fn reset_set_indices(&self) -> Vec<u32> {
-        self.msr_reset
-            .baseline()
-            .iter()
-            .map(|entry| entry.index)
-            .collect()
+        self.msr_reset.reset_indices()
     }
 }
 
@@ -1597,7 +1608,7 @@ mod tests {
         assert_eq!(hyperlight_vm.vm.xcr0().unwrap(), 3);
 
         // Reset the vCPU
-        hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+        hyperlight_vm.reset_vcpu(0, &default_sregs(), &[]).unwrap();
 
         // Verify registers are reset to defaults
         assert_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1761,7 +1772,7 @@ mod tests {
             assert_eq!(regs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vcpu(0, &default_sregs(), &[]).unwrap();
 
             // Check registers are reset to defaults
             assert_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1885,7 +1896,7 @@ mod tests {
             }
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vcpu(0, &default_sregs(), &[]).unwrap();
 
             // Check FPU is reset to defaults
             assert_fpu_reset(hyperlight_vm.vm.as_ref());
@@ -1936,7 +1947,7 @@ mod tests {
             assert_eq!(debug_regs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vcpu(0, &default_sregs(), &[]).unwrap();
 
             // Check debug registers are reset to default values
             assert_debug_regs_reset(hyperlight_vm.vm.as_ref());
@@ -1985,7 +1996,7 @@ mod tests {
             assert_eq!(sregs, expected_dirty);
 
             // Reset vcpu
-            hyperlight_vm.reset_vcpu(0, &default_sregs()).unwrap();
+            hyperlight_vm.reset_vcpu(0, &default_sregs(), &[]).unwrap();
 
             // Check registers are reset to defaults (CR3 is 0 as passed to reset_vcpu)
             let sregs = hyperlight_vm.vm.sregs().unwrap();
@@ -2023,7 +2034,10 @@ mod tests {
             let root_pt_addr = ctx.ctx.vm.get_root_pt().unwrap();
             let segment_state = ctx.ctx.vm.get_snapshot_sregs().unwrap();
 
-            ctx.ctx.vm.reset_vcpu(root_pt_addr, &segment_state).unwrap();
+            ctx.ctx
+                .vm
+                .reset_vcpu(root_pt_addr, &segment_state, &[])
+                .unwrap();
 
             // Re-run from entrypoint (flag=1 means guest skips dirty phase, just does FXSAVE)
             // Use stack_top - 8 to match initialise()'s behavior (simulates call pushing return addr)

src/hyperlight_host/src/hypervisor/regs/x86_64/msrs.rs

@@ -60,9 +60,10 @@ impl MsrResetState {
         })
     }
 
-    /// The creation-time baseline entries.
-    pub fn baseline(&self) -> &[MsrEntry] {
-        &self.baseline
+    /// Every MSR index in the reset set.
+    #[cfg(test)]
+    pub fn reset_indices(&self) -> Vec<u32> {
+        self.baseline.iter().map(|entry| entry.index).collect()
     }
 
     /// The MSR indices captured into a snapshot: the declared guest MSRs plus

src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs

@@ -309,6 +309,12 @@ pub enum RegisterError {
     #[error("Failed to set MSRs: {0}")]
     SetMsrs(HypervisorError),
     #[cfg(target_arch = "x86_64")]
+    #[error("Failed to set batched registers: {0}")]
+    SetBatchedRegisters(HypervisorError),
+    #[cfg(target_arch = "x86_64")]
+    #[error("Batched register writes are not supported")]
+    BatchedSetRegistersUnsupported,
+    #[cfg(target_arch = "x86_64")]
     #[error("Snapshot MSR index {index:#x} is not in this VM's reset set")]
     InvalidSnapshotMsrIndex {
         /// Architectural MSR index supplied by the snapshot.
@@ -503,6 +509,22 @@ pub(crate) trait VirtualMachine: Debug + Send {
     #[cfg(target_arch = "x86_64")]
     fn set_xcr0(&self, value: u64) -> std::result::Result<(), RegisterError>;
 
+    #[cfg(target_arch = "x86_64")]
+    fn can_batch_registers(&self) -> bool {
+        false
+    }
+    #[cfg(target_arch = "x86_64")]
+    fn set_batched_registers(
+        &mut self,
+        _regs: &CommonRegisters,
+        _debug_regs: &CommonDebugRegs,
+        _sregs: &CommonSpecialRegisters,
+        _xcr0: u64,
+        _msrs: &[MsrEntry],
+    ) -> std::result::Result<(), RegisterError> {
+        Err(RegisterError::BatchedSetRegistersUnsupported)
+    }
+
     /// Single-operation vCPU reset
     #[cfg(target_arch = "aarch64")]
     fn can_reset_vcpu(&self) -> bool {