From 4e48f429ebc655900f09532b0a47a998c0782a2e Mon Sep 17 00:00:00 2001
From: Jinank Jain <jinankjain@microsoft.com>
Date: Wed, 7 May 2025 12:46:33 +0000
Subject: [PATCH] vmm: Unify loading of payload for IGVM and non-IGVM

It just simplifies code and improves the code read-ability without much
affecting the boot performance.

Signed-off-by: Jinank Jain <jinankjain@microsoft.com>
---
 vmm/src/vm.rs | 44 ++++++++++++++++++--------------------------
 1 file changed, 18 insertions(+), 26 deletions(-)

diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index a49374eb6..30a49f339 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -509,13 +509,6 @@ impl Vm {
             .validate()
             .map_err(Error::ConfigValidation)?;
 
-        #[cfg(not(feature = "igvm"))]
-        let load_payload_handle = if snapshot.is_none() {
-            Self::load_payload_async(&memory_manager, &config)?
-        } else {
-            None
-        };
-
         info!("Booting VM from config: {:?}", &config);
 
         // Create NUMA nodes based on NumaConfig.
@@ -580,25 +573,6 @@ impl Vm {
             )
             .map_err(Error::CpuManager)?;
 
-        // Loading the igvm file is pushed down here because
-        // igvm parser needs cpu_manager to retrieve cpuid leaf.
-        // For the regular case, we can start loading early, but for
-        // igvm case we have to wait until cpu_manager is created.
-        // Currently, Microsoft Hypervisor does not provide any
-        // Hypervisor specific common cpuid, we need to call get_cpuid_values
-        // per cpuid through cpu_manager.
-        #[cfg(feature = "igvm")]
-        let load_payload_handle = if snapshot.is_none() {
-            Self::load_payload_async(
-                &memory_manager,
-                &config,
-                &cpu_manager,
-                #[cfg(feature = "sev_snp")]
-                sev_snp_enabled,
-            )?
-        } else {
-            None
-        };
         // The initial TDX configuration must be done before the vCPUs are
         // created
         #[cfg(feature = "tdx")]
@@ -672,6 +646,24 @@ impl Vm {
             }
         }
 
+        // Loading the igvm file is pushed down here because
+        // igvm parser needs cpu_manager to retrieve cpuid leaf.
+        // Currently, Microsoft Hypervisor does not provide any
+        // Hypervisor specific common cpuid, we need to call get_cpuid_values
+        // per cpuid through cpu_manager.
+        let load_payload_handle = if snapshot.is_none() {
+            Self::load_payload_async(
+                &memory_manager,
+                &config,
+                #[cfg(feature = "igvm")]
+                &cpu_manager,
+                #[cfg(feature = "sev_snp")]
+                sev_snp_enabled,
+            )?
+        } else {
+            None
+        };
+
         cpu_manager
             .lock()
             .unwrap()