src/: Replace some __PRE_RAM__ use

Change-Id: Iaa56e7b98aad33eeb876edd7465c56c80fd1ac18
Signed-off-by: Kyösti Mälkki <kyosti.malkki@gmail.com>
Reviewed-on: https://review.coreboot.org/c/coreboot/+/35398
Tested-by: build bot (Jenkins) <no-reply@coreboot.org>
Reviewed-by: Aaron Durbin <adurbin@chromium.org>
diff --git a/src/lib/imd_cbmem.c b/src/lib/imd_cbmem.c
index e5c3e41..f04e086 100644
--- a/src/lib/imd_cbmem.c
+++ b/src/lib/imd_cbmem.c
@@ -172,14 +172,13 @@
 	if (imd_recover(imd))
 		return 1;
 
-#if defined(__PRE_RAM__)
 	/*
 	 * Lock the imd in romstage on a recovery. The assumption is that
 	 * if the imd area was recovered in romstage then S3 resume path
 	 * is being taken.
 	 */
-	imd_lockdown(imd);
-#endif
+	if (ENV_ROMSTAGE)
+		imd_lockdown(imd);
 
 	/* Add the specified range first */
 	if (size)