4 files changed, 15 insertions, 8 deletions
diff --git a/src/cpu/x86/smm/smm_module_handler.c b/src/cpu/x86/smm/smm_module_handler.c
index dec2d6b72f..f7cf928f73 100644
--- a/src/cpu/x86/smm/smm_module_handler.c
+++ b/src/cpu/x86/smm/smm_module_handler.c
@@ -94,15 +94,10 @@ struct global_nvs *gnvs;
 
 void *smm_get_save_state(int cpu)
 {
-	char *base;
+	if (cpu > smm_runtime.num_cpus)
+		return NULL;
 
-	/* This function assumes all save states start at top of default
-	 * SMRAM size space and are staggered down by save state size. */
-	base = (void *)(uintptr_t)smm_runtime.smbase;
-	base += SMM_DEFAULT_SIZE;
-	base -= (cpu + 1) * smm_runtime.save_state_size;
-
-	return base;
+	return (void *)(smm_runtime.save_state_top[cpu] - smm_runtime.save_state_size);
 }
 
 uint32_t smm_revision(void)
diff --git a/src/cpu/x86/smm/smm_module_loader.c b/src/cpu/x86/smm/smm_module_loader.c
index de2738ab36..021d739f94 100644
--- a/src/cpu/x86/smm/smm_module_loader.c
+++ b/src/cpu/x86/smm/smm_module_loader.c
@@ -1,5 +1,6 @@
 /* SPDX-License-Identifier: GPL-2.0-only */
 
+#include <stdint.h>
 #include <string.h>
 #include <acpi/acpi_gnvs.h>
 #include <rmodule.h>
@@ -392,5 +393,10 @@ int smm_load_module(void *smram, size_t size, struct smm_loader_params *params)
 	handler_mod_params->num_cpus = params->num_concurrent_stacks;
 	handler_mod_params->gnvs_ptr = (uintptr_t)acpi_get_gnvs();
 
+	for (int i = 0; i < CONFIG_MAX_CPUS; i++) {
+		handler_mod_params->save_state_top[i] = (uintptr_t)smram + SMM_DEFAULT_SIZE
+			- params->per_cpu_save_state_size * i;
+	}
+
 	return smm_module_setup_stub(smram, size, params, fxsave_area);
 }
diff --git a/src/cpu/x86/smm/smm_module_loaderv2.c b/src/cpu/x86/smm/smm_module_loaderv2.c
index e0877e7fc4..886041dbba 100644
--- a/src/cpu/x86/smm/smm_module_loaderv2.c
+++ b/src/cpu/x86/smm/smm_module_loaderv2.c
@@ -641,5 +641,10 @@ int smm_load_module(void *smram, size_t size, struct smm_loader_params *params)
 		return -1;
 	}
 
+	for (int i = 0; i < params->num_concurrent_stacks; i++) {
+		handler_mod_params->save_state_top[i] =
+			cpus[i].ss_start + params->per_cpu_save_state_size;
+	}
+
 	return smm_module_setup_stub(base, size, params, fxsave_area);
 }
diff --git a/src/include/cpu/x86/smm.h b/src/include/cpu/x86/smm.h
index cb9aeac813..b100bf800b 100644
--- a/src/include/cpu/x86/smm.h
+++ b/src/include/cpu/x86/smm.h
@@ -60,6 +60,7 @@ struct smm_runtime {
 	u32 save_state_size;
 	u32 num_cpus;
 	u32 gnvs_ptr;
+	uintptr_t save_state_top[CONFIG_MAX_CPUS];
 } __packed;
 
 struct smm_module_params {