CBMEM: Change declarations for initialization hooks

There are efforts to have bootflows that do not follow a traditional
bootblock-romstage-postcar-ramstage model. As part of that CBMEM
initialisation hooks will need to move from romstage to bootblock.

The interface towards platforms and drivers will change to use one of
CBMEM_CREATION_HOOK() or CBMEM_READY_HOOK(). Former will only be called
in the first stage with CBMEM available.

Change-Id: Ie24bf4e818ca69f539196c3a814f3c52d4103d7e
Signed-off-by: Kyösti Mälkki <kyosti.malkki@gmail.com>
Reviewed-on: https://review.coreboot.org/c/coreboot/+/63375
Tested-by: build bot (Jenkins) <no-reply@coreboot.org>
Reviewed-by: Arthur Heymans <arthur@aheymans.xyz>
Reviewed-by: Julius Werner <jwerner@chromium.org>
Reviewed-by: Raul Rangel <rrangel@chromium.org>
diff --git a/src/drivers/intel/fsp2_0/hand_off_block.c b/src/drivers/intel/fsp2_0/hand_off_block.c
index 9144eaa..e166cf7 100644
--- a/src/drivers/intel/fsp2_0/hand_off_block.c
+++ b/src/drivers/intel/fsp2_0/hand_off_block.c
@@ -109,7 +109,7 @@
 	*cbmem_loc = (uintptr_t)hob_list;
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(save_hob_list);
+CBMEM_CREATION_HOOK(save_hob_list);
 
 const void *fsp_get_hob_list(void)
 {
diff --git a/src/drivers/usb/ehci_debug.c b/src/drivers/usb/ehci_debug.c
index db42ce4..178877b 100644
--- a/src/drivers/usb/ehci_debug.c
+++ b/src/drivers/usb/ehci_debug.c
@@ -680,7 +680,7 @@
 	struct ehci_debug_info *dbg_info_cbmem;
 	int rv;
 
-	if (ENV_ROMSTAGE) {
+	if (ENV_CREATES_CBMEM) {
 		/* Move state from CAR to CBMEM. */
 		struct ehci_debug_info *dbg_info = dbgp_ehci_info();
 		dbg_info_cbmem = cbmem_add(CBMEM_ID_EHCI_DEBUG,
@@ -706,9 +706,7 @@
 		printk(BIOS_DEBUG, "usbdebug: " ENV_STRING " starting...\n");
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_ehci_debug);
-POSTCAR_CBMEM_INIT_HOOK(migrate_ehci_debug);
-RAMSTAGE_CBMEM_INIT_HOOK(migrate_ehci_debug);
+CBMEM_READY_HOOK(migrate_ehci_debug);
 
 int dbgp_ep_is_active(struct dbgp_pipe *pipe)
 {
@@ -728,7 +726,7 @@
 void usbdebug_init(void)
 {
 	/* USB console init is done early in romstage, yet delayed to
-	 * CBMEM_INIT_HOOKs for postcar and ramstage as we recover state
+	 * CBMEM_READY_HOOKs for postcar and ramstage as we recover state
 	 * from CBMEM.
 	 */
 	if (CONFIG(USBDEBUG_IN_PRE_RAM)
diff --git a/src/drivers/vpd/vpd.c b/src/drivers/vpd/vpd.c
index 32ec0a0..24b878f 100644
--- a/src/drivers/vpd/vpd.c
+++ b/src/drivers/vpd/vpd.c
@@ -295,4 +295,4 @@
 	return true;
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(cbmem_add_cros_vpd)
+CBMEM_CREATION_HOOK(cbmem_add_cros_vpd);
diff --git a/src/include/cbmem.h b/src/include/cbmem.h
index cd7751c..5b5191a 100644
--- a/src/include/cbmem.h
+++ b/src/include/cbmem.h
@@ -106,49 +106,35 @@
 void cbmem_list(void);
 void cbmem_add_records_to_cbtable(struct lb_header *header);
 
-#if ENV_RAMSTAGE
-#define ROMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_rom_ = init_fn_;
-#define RAMSTAGE_CBMEM_INIT_HOOK(init_fn_) \
+#define _CBMEM_INIT_HOOK_UNUSED(init_fn_) __attribute__((unused)) \
+	static cbmem_init_hook_t init_fn_ ## _unused_ = init_fn_
+
+#define _CBMEM_INIT_HOOK(init_fn_) \
 	static cbmem_init_hook_t init_fn_ ## _ptr_ __attribute__((used, \
-	section(".rodata.cbmem_init_hooks"))) = init_fn_;
-#define POSTCAR_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_pc_ = init_fn_;
-#elif ENV_ROMSTAGE
-#define ROMSTAGE_CBMEM_INIT_HOOK(init_fn_) \
-	static cbmem_init_hook_t init_fn_ ## _ptr_ __attribute__((used, \
-	section(".rodata.cbmem_init_hooks"))) = init_fn_;
-#define RAMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_ram_ = init_fn_;
-#define POSTCAR_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_pc_ = init_fn_;
-#elif ENV_POSTCAR
-#define ROMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_rom_ = init_fn_;
-#define RAMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_ram_ = init_fn_;
-#define POSTCAR_CBMEM_INIT_HOOK(init_fn_) \
-	static cbmem_init_hook_t init_fn_ ## _ptr_ __attribute__((used, \
-	section(".rodata.cbmem_init_hooks"))) = init_fn_;
-#else
-#define ROMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_rom_ = init_fn_;
-#define RAMSTAGE_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_ram_ = init_fn_;
-#define POSTCAR_CBMEM_INIT_HOOK(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_pc_ = init_fn_;
-#endif /* ENV_RAMSTAGE */
+	section(".rodata.cbmem_init_hooks"))) = init_fn_
 
 /* Early hooks get executed before other hooks. Use sparingly for hooks that create
    CBMEM regions which need to remain in a constant location across boot modes. */
-#if ENV_ROMSTAGE
-#define ROMSTAGE_CBMEM_INIT_HOOK_EARLY(init_fn_) \
+#define _CBMEM_INIT_HOOK_EARLY(init_fn_) \
 	static cbmem_init_hook_t init_fn_ ## _ptr_ __attribute__((used, \
-	section(".rodata.cbmem_init_hooks_early"))) = init_fn_;
+	section(".rodata.cbmem_init_hooks_early"))) = init_fn_
+
+/* Use CBMEM_CREATION_HOOK for code that needs to be run when cbmem is initialized
+   for the first time. */
+#if ENV_CREATES_CBMEM
+#define CBMEM_CREATION_HOOK(x)		_CBMEM_INIT_HOOK(x)
 #else
-#define ROMSTAGE_CBMEM_INIT_HOOK_EARLY(init_fn_) __attribute__((unused)) \
-	static cbmem_init_hook_t init_fn_ ## _unused_rom_ = init_fn_;
-#endif /* ENV_ROMSTAGE */
+#define CBMEM_CREATION_HOOK(x)		_CBMEM_INIT_HOOK_UNUSED(x)
+#endif
+
+/* Use CBMEM_READY_HOOK for code that needs to run in all stages that have cbmem. */
+#if ENV_HAS_CBMEM
+#define CBMEM_READY_HOOK(x)		_CBMEM_INIT_HOOK(x)
+#define CBMEM_READY_HOOK_EARLY(x)	_CBMEM_INIT_HOOK_EARLY(x)
+#else
+#define CBMEM_READY_HOOK(x)		_CBMEM_INIT_HOOK_UNUSED(x)
+#define CBMEM_READY_HOOK_EARLY(x)	_CBMEM_INIT_HOOK_UNUSED(x)
+#endif
 
 /*
  * Returns 0 for the stages where we know that cbmem does not come online.
diff --git a/src/include/rules.h b/src/include/rules.h
index 8ccfe28..4bcd8ea 100644
--- a/src/include/rules.h
+++ b/src/include/rules.h
@@ -292,6 +292,9 @@
 #define ENV_INITIAL_STAGE		ENV_BOOTBLOCK
 #endif
 
+#define ENV_CREATES_CBMEM	ENV_ROMSTAGE
+#define ENV_HAS_CBMEM		(ENV_ROMSTAGE | ENV_POSTCAR | ENV_RAMSTAGE)
+
 #if ENV_X86
 #define ENV_HAS_SPINLOCKS		!ENV_ROMSTAGE_OR_BEFORE
 #elif ENV_RISCV
diff --git a/src/lib/cbfs.c b/src/lib/cbfs.c
index 8c75273..f8112d6 100644
--- a/src/lib/cbfs.c
+++ b/src/lib/cbfs.c
@@ -32,7 +32,7 @@
 		mem_pool_init(&cbfs_cache, _postram_cbfs_cache, REGION_SIZE(postram_cbfs_cache),
 			      CONFIG_CBFS_CACHE_ALIGN);
 }
-ROMSTAGE_CBMEM_INIT_HOOK(switch_to_postram_cache);
+CBMEM_CREATION_HOOK(switch_to_postram_cache);
 
 enum cb_err _cbfs_boot_lookup(const char *name, bool force_ro,
 			      union cbfs_mdata *mdata, struct region_device *rdev)
@@ -693,5 +693,5 @@
 	mcache_to_cbmem(vboot_get_cbfs_boot_device(), CBMEM_ID_CBFS_RW_MCACHE);
 	mcache_to_cbmem(cbfs_get_boot_device(true), CBMEM_ID_CBFS_RO_MCACHE);
 }
-ROMSTAGE_CBMEM_INIT_HOOK(cbfs_mcache_migrate)
+CBMEM_CREATION_HOOK(cbfs_mcache_migrate);
 #endif
diff --git a/src/lib/cbmem_console.c b/src/lib/cbmem_console.c
index 783f336..ad3b99b 100644
--- a/src/lib/cbmem_console.c
+++ b/src/lib/cbmem_console.c
@@ -167,12 +167,10 @@
 	copy_console_buffer(previous_cons_p);
 }
 
-/* Run the romstage hook early so that the console region is one of the earliest created, and
+/* Run this hook early so that the console region is one of the earliest created, and
    therefore more likely to stay in the same place even across different boot modes where some
    other regions may sometimes not get created (e.g. RW_MCACHE in vboot recovery mode). */
-ROMSTAGE_CBMEM_INIT_HOOK_EARLY(cbmemc_reinit)
-RAMSTAGE_CBMEM_INIT_HOOK(cbmemc_reinit)
-POSTCAR_CBMEM_INIT_HOOK(cbmemc_reinit)
+CBMEM_READY_HOOK_EARLY(cbmemc_reinit);
 
 #if CONFIG(CONSOLE_CBMEM_DUMP_TO_UART)
 void cbmem_dump_console_to_uart(void)
diff --git a/src/lib/ext_stage_cache.c b/src/lib/ext_stage_cache.c
index d498597..a35ce37 100644
--- a/src/lib/ext_stage_cache.c
+++ b/src/lib/ext_stage_cache.c
@@ -158,6 +158,4 @@
 		stage_cache_create_empty();
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(stage_cache_setup)
-RAMSTAGE_CBMEM_INIT_HOOK(stage_cache_setup)
-POSTCAR_CBMEM_INIT_HOOK(stage_cache_setup)
+CBMEM_READY_HOOK(stage_cache_setup);
diff --git a/src/lib/fmap.c b/src/lib/fmap.c
index 4d197db..3889dd5 100644
--- a/src/lib/fmap.c
+++ b/src/lib/fmap.c
@@ -315,13 +315,11 @@
 
 static void fmap_setup_cbmem_cache(int unused)
 {
-	if (ENV_ROMSTAGE)
+	if (ENV_CREATES_CBMEM)
 		fmap_add_cbmem_cache();
 
 	/* Finally advertise the cache for the current stage */
 	fmap_register_cbmem_cache();
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(fmap_setup_cbmem_cache)
-RAMSTAGE_CBMEM_INIT_HOOK(fmap_setup_cbmem_cache)
-POSTCAR_CBMEM_INIT_HOOK(fmap_setup_cbmem_cache)
+CBMEM_READY_HOOK(fmap_setup_cbmem_cache);
diff --git a/src/lib/imd_cbmem.c b/src/lib/imd_cbmem.c
index 453b0c1..a855cf1 100644
--- a/src/lib/imd_cbmem.c
+++ b/src/lib/imd_cbmem.c
@@ -17,7 +17,7 @@
 
 void *cbmem_top(void)
 {
-	if (ENV_ROMSTAGE) {
+	if (ENV_CREATES_CBMEM) {
 		static void *top;
 		if (top)
 			return top;
@@ -49,9 +49,8 @@
 
 static void cbmem_top_init_once(void)
 {
-	/* Call one-time hook on expected cbmem init during boot. This sequence
-	   assumes first init call is in romstage. */
-	if (!ENV_ROMSTAGE)
+	/* Call one-time hook on expected cbmem init during boot. */
+	if (!ENV_CREATES_CBMEM)
 		return;
 
 	/* The test is only effective on X86 and when address hits UC memory. */
@@ -106,7 +105,7 @@
 	 * if the imd area was recovered in romstage then S3 resume path
 	 * is being taken.
 	 */
-	if (ENV_ROMSTAGE)
+	if (ENV_CREATES_CBMEM)
 		imd_lockdown(&imd);
 
 	/* Add the specified range first */
@@ -206,8 +205,7 @@
 	imd_region_used(&imd, baseptr, size);
 }
 
-#if ENV_PAYLOAD_LOADER || (CONFIG(EARLY_CBMEM_LIST) \
-	&& (ENV_POSTCAR || ENV_ROMSTAGE))
+#if ENV_PAYLOAD_LOADER || (CONFIG(EARLY_CBMEM_LIST) && ENV_HAS_CBMEM)
 /*
  * -fdata-sections doesn't work so well on read only strings. They all
  * get put in the same section even though those strings may never be
diff --git a/src/lib/program.ld b/src/lib/program.ld
index a6d62cb..8db7ddc 100644
--- a/src/lib/program.ld
+++ b/src/lib/program.ld
@@ -24,7 +24,7 @@
 	*(.text);
 	*(.text.*);
 
-#if ENV_RAMSTAGE || ENV_ROMSTAGE || ENV_POSTCAR
+#if ENV_HAS_CBMEM
 	. = ALIGN(ARCH_POINTER_ALIGN_SIZE);
 	_cbmem_init_hooks = .;
 	KEEP(*(.rodata.cbmem_init_hooks_early));
diff --git a/src/lib/timestamp.c b/src/lib/timestamp.c
index b92975f..c7e1c0c 100644
--- a/src/lib/timestamp.c
+++ b/src/lib/timestamp.c
@@ -215,7 +215,7 @@
 
 	/* First time into romstage we make a clean new table. For platforms that travel
 	   through this path on resume, ARCH_X86 S3, timestamps are also reset. */
-	if (ENV_ROMSTAGE) {
+	if (ENV_CREATES_CBMEM) {
 		ts_cbmem_table = timestamp_alloc_cbmem_table();
 	} else {
 		/* Find existing table in cbmem. */
@@ -228,7 +228,7 @@
 		return;
 	}
 
-	if (ENV_ROMSTAGE)
+	if (ENV_CREATES_CBMEM)
 		timestamp_sync_cache_to_cbmem(ts_cbmem_table);
 
 	/* Seed the timestamp tick frequency in ENV_PAYLOAD_LOADER. */
@@ -279,9 +279,7 @@
 	return (timestamp_get() - ts->base_time) / ts->tick_freq_mhz;
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(timestamp_reinit)
-POSTCAR_CBMEM_INIT_HOOK(timestamp_reinit)
-RAMSTAGE_CBMEM_INIT_HOOK(timestamp_reinit)
+CBMEM_READY_HOOK(timestamp_reinit);
 
 /* Provide default timestamp implementation using monotonic timer. */
 uint64_t  __weak timestamp_get(void)
diff --git a/src/mainboard/google/gale/mmu.c b/src/mainboard/google/gale/mmu.c
index 3b9917c..1d8bd33 100644
--- a/src/mainboard/google/gale/mmu.c
+++ b/src/mainboard/google/gale/mmu.c
@@ -28,7 +28,7 @@
 		/* Map DMA memory */
 		mmu_config_range(DMA_START, DMA_SIZE, DCACHE_OFF);
 		/* Mark cbmem backing store as ready. */
-		if (ENV_ROMSTAGE)
+		if (ENV_CREATES_CBMEM)
 			ipq_cbmem_backing_store_ready();
 	} else {
 		mmu_disable_range(DRAM_START, DRAM_SIZE);
diff --git a/src/mainboard/google/storm/mmu.c b/src/mainboard/google/storm/mmu.c
index 37b4654..b482f9b 100644
--- a/src/mainboard/google/storm/mmu.c
+++ b/src/mainboard/google/storm/mmu.c
@@ -25,7 +25,7 @@
 		mmu_config_range(DRAM_START, DRAM_SIZE, DCACHE_WRITEBACK);
 		/* Map DMA memory */
 		mmu_config_range(DMA_START, DMA_SIZE, DCACHE_OFF);
-		if (ENV_ROMSTAGE)
+		if (ENV_CREATES_CBMEM)
 			/* Mark cbmem backing store as ready. */
 			ipq_cbmem_backing_store_ready();
 	} else {
diff --git a/src/security/tpm/tspi/log.c b/src/security/tpm/tspi/log.c
index 1d6f9ac..296cb2d 100644
--- a/src/security/tpm/tspi/log.c
+++ b/src/security/tpm/tspi/log.c
@@ -39,8 +39,8 @@
 	if (!cbmem_possibly_online() &&
 		!CONFIG(VBOOT_RETURN_FROM_VERSTAGE))
 		return (struct tcpa_table *)_tpm_tcpa_log;
-	else if (ENV_ROMSTAGE &&
-		!CONFIG(VBOOT_RETURN_FROM_VERSTAGE)) {
+	else if (ENV_CREATES_CBMEM
+		 && !CONFIG(VBOOT_RETURN_FROM_VERSTAGE)) {
 		tclt = tcpa_cbmem_init();
 		if (!tclt)
 			return (struct tcpa_table *)_tpm_tcpa_log;
@@ -154,7 +154,7 @@
 		memcpy(tce->digest, preram_log->entries[i].digest, tce->digest_length);
 	}
 }
-ROMSTAGE_CBMEM_INIT_HOOK(recover_tcpa_log);
+CBMEM_CREATION_HOOK(recover_tcpa_log);
 #endif
 
 BOOT_STATE_INIT_ENTRY(BS_PAYLOAD_BOOT, BS_ON_ENTRY, tcpa_log_dump, NULL);
diff --git a/src/security/vboot/common.c b/src/security/vboot/common.c
index f3e2e4b..c2d9340 100644
--- a/src/security/vboot/common.c
+++ b/src/security/vboot/common.c
@@ -96,4 +96,4 @@
 
 	assert(rv == VB2_SUCCESS);
 }
-ROMSTAGE_CBMEM_INIT_HOOK(vboot_setup_cbmem)
+CBMEM_CREATION_HOOK(vboot_setup_cbmem);
diff --git a/src/soc/amd/common/block/pm/chipset_state.c b/src/soc/amd/common/block/pm/chipset_state.c
index 42d1a15..5dd6bd0 100644
--- a/src/soc/amd/common/block/pm/chipset_state.c
+++ b/src/soc/amd/common/block/pm/chipset_state.c
@@ -31,4 +31,4 @@
 		memcpy(state, &chipset_state, sizeof(*state));
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(add_chipset_state_cbmem);
+CBMEM_CREATION_HOOK(add_chipset_state_cbmem);
diff --git a/src/soc/amd/common/pi/heapmanager.c b/src/soc/amd/common/pi/heapmanager.c
index 699bb53..ab923b9 100644
--- a/src/soc/amd/common/pi/heapmanager.c
+++ b/src/soc/amd/common/pi/heapmanager.c
@@ -127,7 +127,7 @@
 	memset(SecondNodePtr, 0, sizeof(BIOS_BUFFER_NODE));
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(EmptyHeap)
+CBMEM_CREATION_HOOK(EmptyHeap);
 
 AGESA_STATUS agesa_AllocateBuffer(uint32_t Func, uintptr_t Data,
 							void *ConfigPtr)
diff --git a/src/soc/amd/stoneyridge/romstage.c b/src/soc/amd/stoneyridge/romstage.c
index b7731b0..be40e21 100644
--- a/src/soc/amd/stoneyridge/romstage.c
+++ b/src/soc/amd/stoneyridge/romstage.c
@@ -202,4 +202,4 @@
 		acpi_pm_gpe_add_events_print_events();
 	}
 }
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_power_state)
+CBMEM_CREATION_HOOK(migrate_power_state);
diff --git a/src/soc/intel/baytrail/romstage/romstage.c b/src/soc/intel/baytrail/romstage/romstage.c
index 74238e2..658da2c 100644
--- a/src/soc/intel/baytrail/romstage/romstage.c
+++ b/src/soc/intel/baytrail/romstage/romstage.c
@@ -32,7 +32,7 @@
 	}
 	memcpy(ps_cbmem, ps_car, sizeof(*ps_cbmem));
 }
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_power_state)
+CBMEM_CREATION_HOOK(migrate_power_state);
 
 static struct chipset_power_state *fill_power_state(void)
 {
diff --git a/src/soc/intel/braswell/romstage/romstage.c b/src/soc/intel/braswell/romstage/romstage.c
index daa1b2d..a1623e1 100644
--- a/src/soc/intel/braswell/romstage/romstage.c
+++ b/src/soc/intel/braswell/romstage/romstage.c
@@ -26,7 +26,7 @@
 	}
 	memcpy(ps_cbmem, &power_state, sizeof(*ps_cbmem));
 }
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_power_state);
+CBMEM_CREATION_HOOK(migrate_power_state);
 
 struct chipset_power_state *fill_power_state(void)
 {
diff --git a/src/soc/intel/broadwell/pch/power_state.c b/src/soc/intel/broadwell/pch/power_state.c
index 12bee7d..b1ab622 100644
--- a/src/soc/intel/broadwell/pch/power_state.c
+++ b/src/soc/intel/broadwell/pch/power_state.c
@@ -30,7 +30,7 @@
 	}
 	memcpy(ps_cbmem, ps_car, sizeof(*ps_cbmem));
 }
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_power_state)
+CBMEM_CREATION_HOOK(migrate_power_state);
 
 /* Return 0, 3, or 5 to indicate the previous sleep state. */
 static int prev_sleep_state(const struct chipset_power_state *ps)
diff --git a/src/soc/intel/common/block/pmc/pmclib.c b/src/soc/intel/common/block/pmc/pmclib.c
index 0f1da72..bfe26c3 100644
--- a/src/soc/intel/common/block/pmc/pmclib.c
+++ b/src/soc/intel/common/block/pmc/pmclib.c
@@ -85,7 +85,7 @@
 	}
 	memcpy(ps_cbmem, &power_state, sizeof(*ps_cbmem));
 }
-ROMSTAGE_CBMEM_INIT_HOOK(migrate_power_state)
+CBMEM_CREATION_HOOK(migrate_power_state);
 
 static void print_num_status_bits(int num_bits, uint32_t status,
 				  const char *const bit_names[])
diff --git a/src/soc/intel/quark/storage_test.c b/src/soc/intel/quark/storage_test.c
index 04c6736..f15d295 100644
--- a/src/soc/intel/quark/storage_test.c
+++ b/src/soc/intel/quark/storage_test.c
@@ -159,7 +159,7 @@
 
 	/* Get the structure addresses */
 	media = NULL;
-	if (ENV_ROMSTAGE)
+	if (ENV_CREATES_CBMEM)
 		media = (struct storage_media *)drivers_storage;
 	else
 		media = cbmem_find(CBMEM_ID_STORAGE_DATA);
@@ -225,7 +225,6 @@
 }
 #endif
 
-#if ENV_ROMSTAGE
 static void copy_storage_structures(int is_recovery)
 {
 	struct storage_media *media;
@@ -242,5 +241,4 @@
 	media->ctrlr = &sdhci_ctrlr->sd_mmc_ctrlr;
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(copy_storage_structures);
-#endif
+CBMEM_CREATION_HOOK(copy_storage_structures);
diff --git a/src/soc/mediatek/common/memory.c b/src/soc/mediatek/common/memory.c
index ab1ccb5..1771313 100644
--- a/src/soc/mediatek/common/memory.c
+++ b/src/soc/mediatek/common/memory.c
@@ -146,7 +146,7 @@
 
 	fill_dram_info(mc, curr_ddr_info);
 }
-ROMSTAGE_CBMEM_INIT_HOOK(add_mem_chip_info);
+CBMEM_CREATION_HOOK(add_mem_chip_info);
 
 static int run_dram_blob(struct dramc_param *dparam)
 {
diff --git a/src/soc/qualcomm/common/qclib.c b/src/soc/qualcomm/common/qclib.c
index b6d04b1..2829cb4 100644
--- a/src/soc/qualcomm/common/qclib.c
+++ b/src/soc/qualcomm/common/qclib.c
@@ -19,7 +19,7 @@
 
 #define QCLIB_VERSION 0
 
-/* store QcLib return data until ROMSTAGE_CBMEM_INIT_HOOK runs */
+/* store QcLib return data until CBMEM_CREATION_HOOK runs */
 static void *mem_chip_addr;
 
 static void write_mem_chip_information(struct qclib_cb_if_table_entry *te)
@@ -51,7 +51,7 @@
 	memcpy(mem_region_base, mem_chip_addr, size);
 }
 
-ROMSTAGE_CBMEM_INIT_HOOK(add_mem_chip_info);
+CBMEM_CREATION_HOOK(add_mem_chip_info);
 
 struct qclib_cb_if_table qclib_cb_if_table = {
 	.magic = QCLIB_MAGIC_NUMBER,