20 files changed, 337 insertions, 367 deletions

diff --git a/src/cpu/intel/haswell/romstage.c b/src/cpu/intel/haswell/romstage.c
index 50150cf51b0a..b8c455bcff42 100644
--- a/src/cpu/intel/haswell/romstage.c
+++ b/src/cpu/intel/haswell/romstage.c
@@ -19,6 +19,7 @@
 
 #include <stdint.h>
 #include <string.h>
+#include <cbfs.h>
 #include <cbmem.h>
 #include <console/console.h>
 #include <arch/cpu.h>
@@ -33,9 +34,9 @@
 #include <device/pci_def.h>
 #include <cpu/x86/lapic.h>
 #include <cbfs.h>
-#include <ramstage_cache.h>
 #include <romstage_handoff.h>
 #include <reset.h>
+#include <stage_cache.h>
 #include <vendorcode/google/chromeos/chromeos.h>
 #if CONFIG_EC_GOOGLE_CHROMEEC
 #include <ec/google/chromeec/ec.h>
@@ -256,13 +257,17 @@ void romstage_common(const struct romstage_params *params)
 
 	if (!wake_from_s3) {
 		cbmem_initialize_empty();
+		stage_cache_create_empty();
 		/* Save data returned from MRC on non-S3 resumes. */
 		save_mrc_data(params->pei_data);
-	} else if (cbmem_initialize()) {
-	#if CONFIG_HAVE_ACPI_RESUME
-		/* Failed S3 resume, reset to come up cleanly */
-		reset_system();
-	#endif
+	} else {
+		stage_cache_recover();
+		if (cbmem_initialize()) {
+		#if CONFIG_HAVE_ACPI_RESUME
+			/* Failed S3 resume, reset to come up cleanly */
+			reset_system();
+		#endif
+		}
 	}
 
 	handoff = romstage_handoff_find_or_add();
@@ -307,18 +312,17 @@ void romstage_after_car(void)
 }
 
 
-#if CONFIG_RELOCATABLE_RAMSTAGE
-#include <ramstage_cache.h>
+#if IS_ENABLED(CONFIG_CACHE_RELOCATED_RAMSTAGE_OUTSIDE_CBMEM)
 
-struct ramstage_cache *ramstage_cache_location(long *size)
+void stage_cache_external_region(void **base, size_t *size)
 {
 	/* The ramstage cache lives in the TSEG region at RESERVED_SMM_OFFSET.
 	 * The top of ram is defined to be the TSEG base address. */
 	*size = RESERVED_SMM_SIZE;
-	return (void *)((uint32_t)cbmem_top() + RESERVED_SMM_OFFSET);
+	*base = (void *)((uint32_t)cbmem_top() + RESERVED_SMM_OFFSET);
 }
 
-void ramstage_cache_invalid(struct ramstage_cache *cache)
+void ramstage_cache_invalid(void)
 {
 #if CONFIG_RESET_ON_INVALID_RAMSTAGE_CACHE
 	reset_system();
diff --git a/src/include/cbfs.h b/src/include/cbfs.h
index 086fa19eb646..05ce2f7f311f 100644
--- a/src/include/cbfs.h
+++ b/src/include/cbfs.h
@@ -2,7 +2,7 @@
  * This file is part of the coreboot project.
  *
  * Copyright (C) 2008 Jordan Crouse <jordan@cosmicpenguin.net>
- * Copyright (C) 2013 The Chromium OS Authors. All rights reserved.
+ * Copyright (C) 2013-2015 Google, Inc.
  *
  * This file is dual-licensed. You can choose between:
  *   - The GNU GPL, version 2, as published by the Free Software Foundation
@@ -90,4 +90,3 @@ void cbfs_set_header_offset(size_t offset);
 static inline void cbfs_set_header_offset(size_t offset) {}
 #endif
 #endif
-
diff --git a/src/include/cbmem.h b/src/include/cbmem.h
index 615510b1fb13..8c61cb2a1818 100644
--- a/src/include/cbmem.h
+++ b/src/include/cbmem.h
@@ -78,6 +78,8 @@
 #define CBMEM_ID_SMBIOS         0x534d4254
 #define CBMEM_ID_SMM_SAVE_SPACE	0x07e9acee
 #define CBMEM_ID_SPINTABLE	0x59175917
+#define CBMEM_ID_STAGEx_META	0x57a9e000
+#define CBMEM_ID_STAGEx_CACHE	0x57a9e100
 #define CBMEM_ID_TIMESTAMP	0x54494d45
 #define CBMEM_ID_VBOOT_HANDOFF	0x780074f0
 #define CBMEM_ID_WIFI_CALIBRATION 0x57494649
diff --git a/src/include/program_loading.h b/src/include/program_loading.h
index 269988777f06..02e6f5fdbb4f 100644
--- a/src/include/program_loading.h
+++ b/src/include/program_loading.h
@@ -123,18 +123,8 @@ void run_romstage(void);
 /* Run ramstage from romstage. */
 void run_ramstage(void);
 
-struct romstage_handoff;
-#if IS_ENABLED(CONFIG_RELOCATABLE_RAMSTAGE)
-/* Cache the loaded ramstage described by prog. */
-void cache_loaded_ramstage(struct romstage_handoff *, struct prog *p);
-/* Load ramstage from cache filling in struct prog. */
-void load_cached_ramstage(struct romstage_handoff *h, struct prog *p);
-#else
-static inline void cache_loaded_ramstage(struct romstage_handoff *h,
-						struct prog *p) {}
-static inline void load_cached_ramstage(struct romstage_handoff *h,
-					struct prog *p) {}
-#endif
+/* Called when the stage cache couldn't load ramstage on resume. */
+void ramstage_cache_invalid(void);
 
 /***********************
  *   PAYLOAD LOADING   *
diff --git a/src/include/ramstage_cache.h b/src/include/ramstage_cache.h
deleted file mode 100644
index 8d9b09517777..000000000000
--- a/src/include/ramstage_cache.h
+++ /dev/null
@@ -1,53 +0,0 @@
-/*
- * This file is part of the coreboot project.
- *
- * Copyright (C) 2013 Google Inc.
- *
- * This program is free software; you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation; version 2 of the License.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software
- * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
- */
-
-#ifndef _RAMSTAGE_CACHE_
-#define _RAMSTAGE_CACHE_
-
-#include <stddef.h>
-#include <stdint.h>
-
-/* This structure is saved along with the relocated ramstage program when
- * CONFIG_RELOCATED_RAMSTAGE is employed.  For x86, it can used to protect
- * the integrity of the ramstage program on S3 resume by saving a copy of
- * the relocated ramstage in SMM space with the assumption that the SMM region
- * cannot be altered from the OS. The magic value just serves as a quick sanity
- * check. */
-
-#define RAMSTAGE_CACHE_MAGIC 0xf3c3a02a
-
-struct ramstage_cache {
-	uint32_t magic;
-	uint32_t entry_point;
-	uint32_t load_address;
-	uint32_t size;
-	char program[0];
-} __attribute__((packed));
-
-/* Chipset/Board function for obtaining cache location and size. */
-struct ramstage_cache *ramstage_cache_location(long *size);
-/* Chipset/Board function called when cache is invalid on resume. */
-void ramstage_cache_invalid(struct ramstage_cache *cache);
-
-static inline int ramstage_cache_is_valid(const struct ramstage_cache *c)
-{
-	return (c != NULL && c->magic == RAMSTAGE_CACHE_MAGIC);
-}
-
-#endif  /* _RAMSTAGE_CACHE_ */
diff --git a/src/include/romstage_handoff.h b/src/include/romstage_handoff.h
index 376b4fd119c8..09934ca4424f 100644
--- a/src/include/romstage_handoff.h
+++ b/src/include/romstage_handoff.h
@@ -36,10 +36,6 @@ struct romstage_handoff {
 	uint8_t s3_resume;
 	uint8_t reboot_required;
 	uint8_t reserved[2];
-	/* The ramstage_entry_point is cached in the stag loading path. This
-	 * cached value can only be utilized when the chipset code properly
-	 * fills in the s3_resume field above. */
-	uint32_t ramstage_entry_point;
 };
 
 #if defined(__ROMSTAGE__)
diff --git a/src/include/stage_cache.h b/src/include/stage_cache.h
new file mode 100644
index 000000000000..bde53307ae0e
--- /dev/null
+++ b/src/include/stage_cache.h
@@ -0,0 +1,49 @@
+/*
+ * This file is part of the coreboot project.
+ *
+ * Copyright 2015 Google Inc.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; version 2 of the License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc.
+ */
+
+#ifndef _STAGE_CACHE_H_
+#define _STAGE_CACHE_H_
+
+#include <stddef.h>
+#include <stdint.h>
+#include <program_loading.h>
+
+enum {
+	STAGE_RAMSTAGE,
+	STAGE_REFCODE,
+};
+
+/* Create an empty stage cache. */
+void stage_cache_create_empty(void);
+/* Recover existing stage cache. */
+void stage_cache_recover(void);
+/* Cache the loaded stage provided according to the parameters. */
+void stage_cache_add(int stage_id, struct prog *stage);
+/* Load the cached stage at given location returning the stage entry point. */
+void stage_cache_load_stage(int stage_id, struct prog *stage);
+/* Fill in parameters for the external stage cache, if utilized. */
+void stage_cache_external_region(void **base, size_t *size);
+
+/* Metadata associated with each stage. */
+struct stage_cache {
+	uint64_t load_addr;
+	uint64_t entry_addr;
+};
+
+#endif /* _STAGE_CACHE_H_ */
diff --git a/src/lib/Makefile.inc b/src/lib/Makefile.inc
index 5c715d7a61bb..043238a26391 100644
--- a/src/lib/Makefile.inc
+++ b/src/lib/Makefile.inc
@@ -111,7 +111,14 @@ romstage-y += hexdump.c
 romstage-$(CONFIG_REG_SCRIPT) += reg_script.c
 ramstage-$(CONFIG_REG_SCRIPT) += reg_script.c
 
-romstage-$(CONFIG_RELOCATABLE_RAMSTAGE) += ramstage_cache.c
+ifeq ($(CONFIG_CACHE_RELOCATED_RAMSTAGE_OUTSIDE_CBMEM),y)
+ramstage-y += ext_stage_cache.c
+romstage-y += ext_stage_cache.c
+else
+ramstage-$(CONFIG_RELOCATABLE_RAMSTAGE) += cbmem_stage_cache.c
+romstage-$(CONFIG_RELOCATABLE_RAMSTAGE) += cbmem_stage_cache.c
+endif
+
 
 smm-y += cbfs.c cbfs_core.c memcmp.c
 smm-$(CONFIG_COMPILER_GCC) += gcc.c
diff --git a/src/lib/cbmem_stage_cache.c b/src/lib/cbmem_stage_cache.c
new file mode 100644
index 000000000000..b9ee14e7be0a
--- /dev/null
+++ b/src/lib/cbmem_stage_cache.c
@@ -0,0 +1,76 @@
+/*
+ * This file is part of the coreboot project.
+ *
+ * Copyright 2015 Google Inc.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; version 2 of the License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc.
+ */
+
+#include <arch/early_variables.h>
+#include <cbmem.h>
+#include <stage_cache.h>
+#include <string.h>
+
+
+/* Provide empty implementations by default. */
+void __attribute__((weak)) stage_cache_create_empty(void) {}
+void __attribute__((weak)) stage_cache_recover(void) {}
+
+/* Stage cache uses cbmem. */
+void stage_cache_add(int stage_id, struct prog *stage)
+{
+	struct stage_cache *meta;
+	void *c;
+
+	meta = cbmem_add(CBMEM_ID_STAGEx_META + stage_id, sizeof(*meta));
+	if (meta == NULL)
+		return;
+	meta->load_addr = (uintptr_t)prog_start(stage);
+	meta->entry_addr = (uintptr_t)prog_entry(stage);
+
+	c = cbmem_add(CBMEM_ID_STAGEx_CACHE + stage_id, prog_size(stage));
+	if (c == NULL)
+		return;
+
+	memcpy(c, prog_start(stage), prog_size(stage));
+}
+
+void stage_cache_load_stage(int stage_id, struct prog *stage)
+{
+	struct stage_cache *meta;
+	const struct cbmem_entry *e;
+	void *c;
+	size_t size;
+	void *load_addr;
+
+	prog_set_entry(stage, NULL, NULL);
+
+	meta = cbmem_find(CBMEM_ID_STAGEx_META + stage_id);
+	if (meta == NULL)
+		return;
+
+	e = cbmem_entry_find(CBMEM_ID_STAGEx_CACHE + stage_id);
+
+	if (e == NULL)
+		return;
+
+	c = cbmem_entry_start(e);
+	size = cbmem_entry_size(e);
+	load_addr = (void *)(uintptr_t)meta->load_addr;
+
+	memcpy(load_addr, c, size);
+
+	prog_set_area(stage, load_addr, size);
+	prog_set_entry(stage, (void *)(uintptr_t)meta->entry_addr, NULL);
+}
diff --git a/src/lib/ext_stage_cache.c b/src/lib/ext_stage_cache.c
new file mode 100644
index 000000000000..379b9fcffabb
--- /dev/null
+++ b/src/lib/ext_stage_cache.c
@@ -0,0 +1,129 @@
+/*
+ * This file is part of the coreboot project.
+ *
+ * Copyright 2015 Google Inc.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; version 2 of the License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc.
+ */
+
+#include <arch/early_variables.h>
+#include <bootstate.h>
+#include <cbmem.h>
+#include <console/console.h>
+#include <imd.h>
+#include <rules.h>
+#include <stage_cache.h>
+#include <string.h>
+
+static struct imd imd_stage_cache CAR_GLOBAL = { };
+
+static inline struct imd *imd_get(void)
+{
+	return car_get_var_ptr(&imd_stage_cache);
+}
+
+void stage_cache_create_empty(void)
+{
+	struct imd *imd;
+	void *base;
+	size_t size;
+
+	imd = imd_get();
+	stage_cache_external_region(&base, &size);
+	imd_handle_init(imd, (void *)(size + (uintptr_t)base));
+
+	printk(BIOS_DEBUG, "External stage cache:\n");
+	imd_create_tiered_empty(imd, 4096, 4096, 1024, 32);
+	if (imd_limit_size(imd, size))
+		printk(BIOS_DEBUG, "Could not limit stage cache size.\n");
+}
+
+void stage_cache_recover(void)
+{
+	struct imd *imd;
+	void *base;
+	size_t size;
+
+	imd = imd_get();
+	stage_cache_external_region(&base, &size);
+	imd_handle_init(imd, (void *)(size + (uintptr_t)base));
+	if (imd_recover(imd))
+		printk(BIOS_DEBUG, "Unable to recover external stage cache.\n");
+}
+
+void stage_cache_add(int stage_id, struct prog *stage)
+{
+	struct imd *imd;
+	const struct imd_entry *e;
+	struct stage_cache *meta;
+	void *c;
+
+	imd = imd_get();
+	e = imd_entry_add(imd, CBMEM_ID_STAGEx_META + stage_id, sizeof(*meta));
+
+	if (e == NULL)
+		return;
+
+	meta = imd_entry_at(imd, e);
+
+	meta->load_addr = (uintptr_t)prog_start(stage);
+	meta->entry_addr = (uintptr_t)prog_entry(stage);
+
+	e = imd_entry_add(imd, CBMEM_ID_STAGEx_CACHE + stage_id,
+				prog_size(stage));
+
+	if (e == NULL)
+		return;
+
+	c = imd_entry_at(imd, e);
+
+	memcpy(c, prog_start(stage), prog_size(stage));
+}
+
+void stage_cache_load_stage(int stage_id, struct prog *stage)
+{
+	struct imd *imd;
+	struct stage_cache *meta;
+	const struct imd_entry *e;
+	void *c;
+	size_t size;
+
+	imd = imd_get();
+	e = imd_entry_find(imd, CBMEM_ID_STAGEx_META + stage_id);
+	if (e == NULL)
+		return;
+
+	meta = imd_entry_at(imd, e);
+
+	e = imd_entry_find(imd, CBMEM_ID_STAGEx_CACHE + stage_id);
+
+	if (e == NULL)
+		return;
+
+	c = imd_entry_at(imd, e);
+	size = imd_entry_size(imd, e);
+
+	memcpy((void *)(uintptr_t)meta->load_addr, c, size);
+
+	prog_set_area(stage, (void *)(uintptr_t)meta->load_addr, size);
+	prog_set_entry(stage, (void *)(uintptr_t)meta->entry_addr, NULL);
+}
+
+#if ENV_RAMSTAGE
+static void recover_sc(void *unused)
+{
+	stage_cache_recover();
+}
+BOOT_STATE_INIT_ENTRY(BS_PRE_DEVICE, BS_ON_ENTRY, recover_sc, NULL);
+#endif
diff --git a/src/lib/loaders/load_and_run_ramstage.c b/src/lib/loaders/load_and_run_ramstage.c
index b3728a17cfa9..153e38e11e48 100644
--- a/src/lib/loaders/load_and_run_ramstage.c
+++ b/src/lib/loaders/load_and_run_ramstage.c
@@ -23,6 +23,7 @@
 #include <cbfs.h>
 #include <program_loading.h>
 #include <romstage_handoff.h>
+#include <stage_cache.h>
 #include <timestamp.h>
 
 extern const struct prog_loader_ops cbfs_ramstage_loader;
@@ -35,16 +36,20 @@ static const struct prog_loader_ops *loaders[] = {
 	&cbfs_ramstage_loader,
 };
 
-static void
-load_ramstage(const struct prog_loader_ops *ops,
-		struct romstage_handoff *handoff, struct prog *ramstage)
+void __attribute__((weak)) stage_cache_add(int stage_id, struct prog *stage) {}
+void __attribute__((weak)) stage_cache_load_stage(int stage_id,
+							struct prog *stage) {}
+void __attribute__((weak)) ramstage_cache_invalid(void) {}
+
+static void load_ramstage(const struct prog_loader_ops *ops,
+				struct prog *ramstage)
 {
 	timestamp_add_now(TS_START_COPYRAM);
 
 	if (ops->prepare(ramstage))
 		return;
 
-	cache_loaded_ramstage(handoff, ramstage);
+	stage_cache_add(STAGE_RAMSTAGE, ramstage);
 
 	timestamp_add_now(TS_END_COPYRAM);
 
@@ -56,18 +61,18 @@ static void run_ramstage_from_resume(struct romstage_handoff *handoff,
 {
 	if (handoff != NULL && handoff->s3_resume) {
 		/* Load the cached ramstage to runtime location. */
-		load_cached_ramstage(handoff, ramstage);
+		stage_cache_load_stage(STAGE_RAMSTAGE, ramstage);
 
 		if (prog_entry(ramstage) != NULL) {
 			printk(BIOS_DEBUG, "Jumping to image.\n");
 			prog_run(ramstage);
 		}
+		ramstage_cache_invalid();
 	}
 }
 
 void run_ramstage(void)
 {
-	struct romstage_handoff *handoff;
 	const struct prog_loader_ops *ops;
 	int i;
 	struct prog ramstage = {
@@ -75,14 +80,12 @@ void run_ramstage(void)
 		.type = PROG_RAMSTAGE,
 	};
 
-	handoff = romstage_handoff_find_or_add();
-
-	run_ramstage_from_resume(handoff, &ramstage);
+	run_ramstage_from_resume(romstage_handoff_find_or_add(), &ramstage);
 
 	for (i = 0; i < ARRAY_SIZE(loaders); i++) {
 		ops = loaders[i];
 		printk(BIOS_DEBUG, "Trying %s ramstage loader.\n", ops->name);
-		load_ramstage(ops, handoff, &ramstage);
+		load_ramstage(ops, &ramstage);
 	}
 
 	die("Ramstage was not loaded!\n");
diff --git a/src/lib/ramstage_cache.c b/src/lib/ramstage_cache.c
deleted file mode 100644
index 1dd40c7d5e40..000000000000
--- a/src/lib/ramstage_cache.c
+++ /dev/null
@@ -1,149 +0,0 @@
-/*
- * This file is part of the coreboot project.
- *
- * Copyright (C) 2013 Google Inc.
- *
- * This program is free software; you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation; version 2 of the License.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software
- * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
- */
-
-#include <stddef.h>
-#include <string.h>
-#include <console/console.h>
-#include <program_loading.h>
-#include <ramstage_cache.h>
-#include <romstage_handoff.h>
-
-#if CONFIG_CACHE_RELOCATED_RAMSTAGE_OUTSIDE_CBMEM
-
-void cache_loaded_ramstage(struct romstage_handoff *handoff,
-				struct prog *ramstage)
-{
-	struct ramstage_cache *cache;
-	uint32_t total_size;
-	uint32_t ramstage_size;
-	void *ramstage_base;
-	long cache_size = 0;
-
-	ramstage_size = prog_size(ramstage);
-	ramstage_base = prog_start(ramstage);
-
-	cache = ramstage_cache_location(&cache_size);
-
-	if (cache == NULL) {
-		printk(BIOS_DEBUG, "No ramstage cache location provided.\n");
-		return;
-	}
-
-	total_size = sizeof(*cache) + ramstage_size;
-	if (total_size > cache_size) {
-		printk(BIOS_DEBUG, "cache size too small: 0x%08x > 0x%08lx\n",
-		       total_size, cache_size);
-		/* Nuke whatever may be there now just in case. */
-		cache->magic = ~RAMSTAGE_CACHE_MAGIC;
-		return;
-	}
-
-	cache->magic = RAMSTAGE_CACHE_MAGIC;
-	cache->entry_point = (uint32_t)prog_entry(ramstage);
-	cache->load_address = (uint32_t)ramstage_base;
-	cache->size = ramstage_size;
-
-	printk(BIOS_DEBUG, "Saving ramstage to %p.\n", cache);
-
-	/* Copy over the program. */
-	memcpy(&cache->program[0], ramstage_base, ramstage_size);
-
-	if (handoff == NULL)
-		return;
-
-	handoff->ramstage_entry_point = cache->entry_point;
-}
-
-void load_cached_ramstage(struct romstage_handoff *handoff,
-				struct prog *ramstage)
-{
-	struct ramstage_cache *cache;
-	long size = 0;
-
-	cache = ramstage_cache_location(&size);
-
-	if (!ramstage_cache_is_valid(cache)) {
-		printk(BIOS_DEBUG, "Invalid ramstage cache found.\n");
-		ramstage_cache_invalid(cache);
-		return;
-	}
-
-	printk(BIOS_DEBUG, "Loading ramstage from %p.\n", cache);
-
-	prog_set_area(ramstage, (void *)cache->load_address, cache->size);
-	prog_set_entry(ramstage, (void *)cache->entry_point, NULL);
-
-	memcpy((void *)cache->load_address, &cache->program[0], cache->size);
-}
-
-#else
-
-/* Cache relocated ramstage in CBMEM. */
-
-void cache_loaded_ramstage(struct romstage_handoff *handoff,
-				struct prog *ramstage)
-{
-	uint32_t ramstage_size;
-	const struct cbmem_entry *entry;
-
-	if (handoff == NULL)
-		return;
-
-	ramstage_size = prog_size(ramstage);
-	/* cbmem_entry_add() does a find() before add(). */
-	entry = cbmem_entry_add(CBMEM_ID_RAMSTAGE_CACHE, ramstage_size);
-
-	if (entry == NULL)
-		return;
-
-	/* Keep track of the entry point in the handoff structure. */
-	handoff->ramstage_entry_point = (uint32_t)prog_entry(ramstage);
-
-	memcpy(cbmem_entry_start(entry), prog_start(ramstage), ramstage_size);
-}
-
-void load_cached_ramstage(struct romstage_handoff *handoff,
-				struct prog *ramstage)
-{
-	const struct cbmem_entry *entry_cache;
-	const struct cbmem_entry *entry_dest;
-
-	if (handoff == NULL)
-		return;
-
-	entry_cache = cbmem_entry_find(CBMEM_ID_RAMSTAGE_CACHE);
-
-	if (entry_cache == NULL)
-		return;
-
-	entry_dest = cbmem_entry_find(CBMEM_ID_RAMSTAGE);
-
-	if (entry_dest == NULL)
-		return;
-
-	prog_set_area(ramstage, cbmem_entry_start(entry_dest),
-			cbmem_entry_size(entry_dest));
-	prog_set_entry(ramstage, (void *)handoff->ramstage_entry_point, NULL);
-
-	/* Load the cached ramstage copy into the to-be-run region. */
-	memcpy(prog_start(ramstage), cbmem_entry_start(entry_cache),
-	       prog_size(ramstage));
-}
-
-#endif
diff --git a/src/soc/intel/baytrail/refcode.c b/src/soc/intel/baytrail/refcode.c
index b3e612ad95ce..558db3cc8de8 100644
--- a/src/soc/intel/baytrail/refcode.c
+++ b/src/soc/intel/baytrail/refcode.c
@@ -25,7 +25,7 @@
 #include <cpu/x86/tsc.h>
 #include <program_loading.h>
 #include <rmodule.h>
-#include <ramstage_cache.h>
+#include <stage_cache.h>
 #if IS_ENABLED(CONFIG_CHROMEOS)
 #include <vendorcode/google/chromeos/vboot_handoff.h>
 #endif
@@ -33,11 +33,6 @@
 #include <soc/ramstage.h>
 #include <soc/efi_wrapper.h>
 
-static inline struct ramstage_cache *next_cache(struct ramstage_cache *c)
-{
-	return (struct ramstage_cache *)&c->program[c->size];
-}
-
 static void ABI_X86 send_to_console(unsigned char b)
 {
 	console_tx_byte(b);
@@ -45,60 +40,18 @@ static void ABI_X86 send_to_console(unsigned char b)
 
 static efi_wrapper_entry_t load_refcode_from_cache(void)
 {
-	struct ramstage_cache *c;
-	long cache_size;
+	struct prog refcode;
 
 	printk(BIOS_DEBUG, "refcode loading from cache.\n");
 
-	c = ramstage_cache_location(&cache_size);
-
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "Invalid ramstage cache descriptor.\n");
-		return NULL;
-	}
+	stage_cache_load_stage(STAGE_REFCODE, &refcode);
 
-	c = next_cache(c);
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "Invalid refcode cache descriptor.\n");
-		return NULL;
-	}
-
-	printk(BIOS_DEBUG, "Loading cached reference code from 0x%08x(%x)\n",
-	       c->load_address, c->size);
-	memcpy((void *)c->load_address, &c->program[0], c->size);
-
-	return (efi_wrapper_entry_t)c->entry_point;
+	return (efi_wrapper_entry_t)prog_entry(&refcode);
 }
 
 static void cache_refcode(const struct rmod_stage_load *rsl)
 {
-	struct ramstage_cache *c;
-	long cache_size;
-
-	c = ramstage_cache_location(&cache_size);
-
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "No place to cache reference code.\n");
-		return;
-	}
-
-	/* Determine how much remaining cache available. */
-	cache_size -= c->size + sizeof(*c);
-
-	if (cache_size < (sizeof(*c) + prog_size(rsl->prog))) {
-		printk(BIOS_DEBUG, "Not enough cache space for ref code.\n");
-		return;
-	}
-
-	c = next_cache(c);
-	c->magic = RAMSTAGE_CACHE_MAGIC;
-	c->entry_point = (uint32_t)(uintptr_t)prog_entry(rsl->prog);
-	c->load_address = (uint32_t)(uintptr_t)prog_start(rsl->prog);
-	c->size = prog_size(rsl->prog);
-
-	printk(BIOS_DEBUG, "Caching refcode at 0x%p(%x)\n",
-	       &c->program[0], c->size);
-	memcpy(&c->program[0], (void *)c->load_address, c->size);
+	stage_cache_add(STAGE_REFCODE, rsl->prog);
 }
 
 #if IS_ENABLED(CONFIG_CHROMEOS)
diff --git a/src/soc/intel/baytrail/romstage/raminit.c b/src/soc/intel/baytrail/romstage/raminit.c
index fbd38df13d2c..61c122cd1df2 100644
--- a/src/soc/intel/baytrail/romstage/raminit.c
+++ b/src/soc/intel/baytrail/romstage/raminit.c
@@ -25,6 +25,7 @@
 #include <console/console.h>
 #include <device/pci_def.h>
 #include <halt.h>
+#include <stage_cache.h>
 #include <soc/gpio.h>
 #include <soc/intel/common/mrc_cache.h>
 #include <soc/iomap.h>
@@ -169,12 +170,16 @@ void raminit(struct mrc_params *mp, int prev_sleep_state)
 
 	if (prev_sleep_state != 3) {
 		cbmem_initialize_empty();
-	} else if (cbmem_initialize()) {
-	#if CONFIG_HAVE_ACPI_RESUME
-		printk(BIOS_DEBUG, "Failed to recover CBMEM in S3 resume.\n");
-		/* Failed S3 resume, reset to come up cleanly */
-		reset_system();
-	#endif
+		stage_cache_create_empty();
+	} else {
+		stage_cache_recover();
+		if (cbmem_initialize()) {
+		#if CONFIG_HAVE_ACPI_RESUME
+			printk(BIOS_DEBUG, "Failed to recover CBMEM in S3 resume.\n");
+			/* Failed S3 resume, reset to come up cleanly */
+			reset_system();
+		#endif
+		}
 	}
 
 	printk(BIOS_DEBUG, "MRC Wrapper returned %d\n", ret);
diff --git a/src/soc/intel/baytrail/romstage/romstage.c b/src/soc/intel/baytrail/romstage/romstage.c
index 029fee61ee23..9c8bbc40cef7 100644
--- a/src/soc/intel/baytrail/romstage/romstage.c
+++ b/src/soc/intel/baytrail/romstage/romstage.c
@@ -20,18 +20,18 @@
 #include <stddef.h>
 #include <arch/cpu.h>
 #include <arch/io.h>
-#include <arch/cbfs.h>
 #include <arch/stages.h>
 #include <arch/early_variables.h>
 #include <console/console.h>
+#include <cbfs.h>
 #include <cbmem.h>
 #include <cpu/x86/mtrr.h>
 #if CONFIG_EC_GOOGLE_CHROMEEC
 #include <ec/google/chromeec/ec.h>
 #endif
 #include <elog.h>
-#include <ramstage_cache.h>
 #include <romstage_handoff.h>
+#include <stage_cache.h>
 #include <timestamp.h>
 #include <vendorcode/google/chromeos/chromeos.h>
 #include <soc/gpio.h>
@@ -359,7 +359,7 @@ static void *setup_stack_and_mttrs(void)
 	return slot;
 }
 
-void ramstage_cache_invalid(struct ramstage_cache *cache)
+void ramstage_cache_invalid(void)
 {
 #if CONFIG_RESET_ON_INVALID_RAMSTAGE_CACHE
 	/* Perform cold reset on invalid ramstage cache. */
diff --git a/src/soc/intel/baytrail/stage_cache.c b/src/soc/intel/baytrail/stage_cache.c
index d51746fbd193..a4f510d2f4ce 100644
--- a/src/soc/intel/baytrail/stage_cache.c
+++ b/src/soc/intel/baytrail/stage_cache.c
@@ -18,10 +18,10 @@
  */
 
 #include <cbmem.h>
-#include <ramstage_cache.h>
+#include <stage_cache.h>
 #include <soc/smm.h>
 
-struct ramstage_cache *ramstage_cache_location(long *size)
+void stage_cache_external_region(void **base, size_t *size)
 {
 	char *smm_base;
 	/* 1MiB cache size */
@@ -31,5 +31,5 @@ struct ramstage_cache *ramstage_cache_location(long *size)
 	 * cbmem_top(). */
 	smm_base = cbmem_top();
 	*size = cache_size;
-	return (void *)&smm_base[smm_region_size() - cache_size];
+	*base = &smm_base[smm_region_size() - cache_size];
 }
diff --git a/src/soc/intel/broadwell/refcode.c b/src/soc/intel/broadwell/refcode.c
index 6dec3b915805..201825a93b21 100644
--- a/src/soc/intel/broadwell/refcode.c
+++ b/src/soc/intel/broadwell/refcode.c
@@ -25,7 +25,7 @@
 #include <cpu/x86/tsc.h>
 #include <program_loading.h>
 #include <rmodule.h>
-#include <ramstage_cache.h>
+#include <stage_cache.h>
 #include <string.h>
 #if IS_ENABLED(CONFIG_CHROMEOS)
 #include <vendorcode/google/chromeos/vboot_handoff.h>
@@ -34,67 +34,20 @@
 #include <soc/pei_wrapper.h>
 #include <soc/ramstage.h>
 
-static inline struct ramstage_cache *next_cache(struct ramstage_cache *c)
-{
-	return (struct ramstage_cache *)&c->program[c->size];
-}
-
 static pei_wrapper_entry_t load_refcode_from_cache(void)
 {
-	struct ramstage_cache *c;
-	long cache_size;
+	struct prog refcode;
 
 	printk(BIOS_DEBUG, "refcode loading from cache.\n");
 
-	c = ramstage_cache_location(&cache_size);
-
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "Invalid ramstage cache descriptor.\n");
-		return NULL;
-	}
+	stage_cache_load_stage(STAGE_REFCODE, &refcode);
 
-	c = next_cache(c);
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "Invalid refcode cache descriptor.\n");
-		return NULL;
-	}
-
-	printk(BIOS_DEBUG, "Loading cached reference code from 0x%08x(%x)\n",
-	       c->load_address, c->size);
-	memcpy((void *)c->load_address, &c->program[0], c->size);
-
-	return (pei_wrapper_entry_t)c->entry_point;
+	return (pei_wrapper_entry_t)prog_entry(&refcode);
 }
 
 static void cache_refcode(const struct rmod_stage_load *rsl)
 {
-	struct ramstage_cache *c;
-	long cache_size;
-
-	c = ramstage_cache_location(&cache_size);
-
-	if (!ramstage_cache_is_valid(c)) {
-		printk(BIOS_DEBUG, "No place to cache reference code.\n");
-		return;
-	}
-
-	/* Determine how much remaining cache available. */
-	cache_size -= c->size + sizeof(*c);
-
-	if (cache_size < (sizeof(*c) + prog_size(rsl->prog))) {
-		printk(BIOS_DEBUG, "Not enough cache space for ref code.\n");
-		return;
-	}
-
-	c = next_cache(c);
-	c->magic = RAMSTAGE_CACHE_MAGIC;
-	c->entry_point = (uint32_t)(uintptr_t)prog_entry(rsl->prog);
-	c->load_address = (uint32_t)(uintptr_t)prog_start(rsl->prog);
-	c->size = prog_size(rsl->prog);
-
-	printk(BIOS_DEBUG, "Caching refcode at 0x%p(%x)\n",
-	       &c->program[0], c->size);
-	memcpy(&c->program[0], (void *)c->load_address, c->size);
+	stage_cache_add(STAGE_REFCODE, rsl->prog);
 }
 
 #if IS_ENABLED(CONFIG_CHROMEOS)
diff --git a/src/soc/intel/broadwell/romstage/raminit.c b/src/soc/intel/broadwell/romstage/raminit.c
index 1f62fd0c8c17..52ea49100dff 100644
--- a/src/soc/intel/broadwell/romstage/raminit.c
+++ b/src/soc/intel/broadwell/romstage/raminit.c
@@ -29,6 +29,7 @@
 #include <ec/google/chromeec/ec.h>
 #include <ec/google/chromeec/ec_commands.h>
 #endif
+#include <stage_cache.h>
 #include <vendorcode/google/chromeos/chromeos.h>
 #include <soc/intel/common/mrc_cache.h>
 #include <soc/iomap.h>
@@ -111,12 +112,16 @@ void raminit(struct pei_data *pei_data)
 
 	if (pei_data->boot_mode != SLEEP_STATE_S3) {
 		cbmem_initialize_empty();
-	} else if (cbmem_initialize()) {
+		stage_cache_create_empty();
+	} else {
+		stage_cache_recover();
+		if (cbmem_initialize()) {
 #if CONFIG_HAVE_ACPI_RESUME
-		printk(BIOS_DEBUG, "Failed to recover CBMEM in S3 resume.\n");
-		/* Failed S3 resume, reset to come up cleanly */
-		reset_system();
+			printk(BIOS_DEBUG, "Failed to recover CBMEM in S3 resume.\n");
+			/* Failed S3 resume, reset to come up cleanly */
+			reset_system();
 #endif
+		}
 	}
 
 	printk(BIOS_DEBUG, "MRC data at %p %d bytes\n", pei_data->data_to_save,
diff --git a/src/soc/intel/broadwell/romstage/romstage.c b/src/soc/intel/broadwell/romstage/romstage.c
index bf9e2ed3e9f9..1741ae6b827a 100644
--- a/src/soc/intel/broadwell/romstage/romstage.c
+++ b/src/soc/intel/broadwell/romstage/romstage.c
@@ -25,11 +25,12 @@
 #include <arch/stages.h>
 #include <arch/early_variables.h>
 #include <console/console.h>
+#include <cbfs.h>
 #include <cbmem.h>
 #include <cpu/x86/mtrr.h>
 #include <elog.h>
-#include <ramstage_cache.h>
 #include <romstage_handoff.h>
+#include <stage_cache.h>
 #include <timestamp.h>
 #include <soc/me.h>
 #include <soc/pei_data.h>
@@ -142,7 +143,7 @@ void asmlinkage romstage_after_car(void)
 	while (1);
 }
 
-void ramstage_cache_invalid(struct ramstage_cache *cache)
+void ramstage_cache_invalid(void)
 {
 #if CONFIG_RESET_ON_INVALID_RAMSTAGE_CACHE
 	/* Perform cold reset on invalid ramstage cache. */
diff --git a/src/soc/intel/broadwell/stage_cache.c b/src/soc/intel/broadwell/stage_cache.c
index cb34297cdd85..346e673cf50c 100644
--- a/src/soc/intel/broadwell/stage_cache.c
+++ b/src/soc/intel/broadwell/stage_cache.c
@@ -18,11 +18,11 @@
  */
 
 #include <cbmem.h>
-#include <ramstage_cache.h>
 #include <soc/smm.h>
+#include <stage_cache.h>
 #include <stdint.h>
 
-struct ramstage_cache *ramstage_cache_location(long *size)
+void stage_cache_external_region(void **base, size_t *size)
 {
 	/* The ramstage cache lives in the TSEG region.
 	 * The top of ram is defined to be the TSEG base address. */
@@ -30,6 +30,6 @@ struct ramstage_cache *ramstage_cache_location(long *size)
 	offset -= CONFIG_IED_REGION_SIZE;
 	offset -= CONFIG_SMM_RESERVED_SIZE;
 
+	*base = (void *)(cbmem_top() + offset);
 	*size = CONFIG_SMM_RESERVED_SIZE;
-	return (void *)(cbmem_top() + offset);
 }