9 files changed, 65 insertions, 13 deletions
diff --git a/src/arch/arm/Makefile.inc b/src/arch/arm/Makefile.inc
index e739d0b8d5..ffbc8413a7 100644
--- a/src/arch/arm/Makefile.inc
+++ b/src/arch/arm/Makefile.inc
@@ -54,6 +54,7 @@ ifeq ($(CONFIG_ARCH_BOOTBLOCK_ARM),y)
 bootblock-y += id.S
 $(obj)/arch/arm/id.bootblock.o: $(obj)/build.h
 
+bootblock-y += boot.c
 bootblock-y += stages.c
 bootblock-y += eabi_compat.c
 bootblock-y += memset.S
@@ -89,6 +90,7 @@ verstage-y += stages.c
 
 ifeq ($(CONFIG_ARCH_ROMSTAGE_ARM),y)
 
+romstage-y += boot.c
 romstage-y += stages.c
 romstage-y += div0.c
 romstage-y += eabi_compat.c
diff --git a/src/arch/arm/boot.c b/src/arch/arm/boot.c
index 9dc6dbd3cf..6a2934b86e 100644
--- a/src/arch/arm/boot.c
+++ b/src/arch/arm/boot.c
@@ -20,12 +20,17 @@
 #include <arch/cache.h>
 #include <program_loading.h>
 
-void arch_payload_run(const struct payload *payload)
+void arch_prog_run(struct prog *prog)
 {
 	void (*doit)(void *);
 
 	cache_sync_instructions();
 
-	doit = prog_entry(&payload->prog);
-	doit(prog_entry_arg(&payload->prog));
+	doit = prog_entry(prog);
+	doit(prog_entry_arg(prog));
+}
+
+void arch_payload_run(struct payload *payload)
+{
+	arch_prog_run(&payload->prog);
 }
diff --git a/src/arch/arm64/Makefile.inc b/src/arch/arm64/Makefile.inc
index 6d23dab561..c1f1112280 100644
--- a/src/arch/arm64/Makefile.inc
+++ b/src/arch/arm64/Makefile.inc
@@ -57,6 +57,7 @@ bootblock-y += div0.c
 bootblock-y += id.S
 $(obj)/arch/arm64/id.bootblock.o: $(obj)/build.h
 
+bootblock-y += boot.c
 bootblock-y += c_entry.c
 bootblock-y += stage_entry.S
 bootblock-y += stages.c
@@ -81,6 +82,7 @@ endif # CONFIG_ARCH_BOOTBLOCK_ARM64
 
 ifeq ($(CONFIG_ARCH_ROMSTAGE_ARM64),y)
 
+romstage-y += boot.c
 romstage-y += c_entry.c
 romstage-y += stage_entry.S
 romstage-y += stages.c
diff --git a/src/arch/arm64/boot.c b/src/arch/arm64/boot.c
index 74282a9cde..b5e375193e 100644
--- a/src/arch/arm64/boot.c
+++ b/src/arch/arm64/boot.c
@@ -25,15 +25,16 @@
 #include <arch/transition.h>
 #include <console/console.h>
 #include <program_loading.h>
+#include <rules.h>
 #include <string.h>
 
-void arch_payload_run(const struct payload *payload)
+static void run_payload(struct prog *prog)
 {
 	void (*doit)(void *);
 	void *arg;
 
-	doit = prog_entry(&payload->prog);
-	arg = prog_entry_arg(&payload->prog);
+	doit = prog_entry(prog);
+	arg = prog_entry_arg(prog);
 
 	uint8_t current_el = get_current_el();
 
@@ -61,3 +62,23 @@ void arch_payload_run(const struct payload *payload)
 	cache_sync_instructions();
 	transition_with_entry(doit, arg, &exc_state);
 }
+
+void arch_prog_run(struct prog *prog)
+{
+	void (*doit)(void *);
+	void *arg;
+
+	if (ENV_RAMSTAGE && prog->type == PROG_PAYLOAD)
+		run_payload(prog);
+		return;
+
+	doit = prog_entry(prog);
+	arg = prog_entry_arg(prog);
+
+	doit(prog_entry_arg(prog));
+}
+
+void arch_payload_run(struct payload *payload)
+{
+	arch_prog_run(&payload->prog);
+}
diff --git a/src/arch/mips/boot.c b/src/arch/mips/boot.c
index 6e12d10c71..ebc45bb682 100644
--- a/src/arch/mips/boot.c
+++ b/src/arch/mips/boot.c
@@ -20,7 +20,12 @@
 #include <arch/stages.h>
 #include <program_loading.h>
 
-void arch_payload_run(const struct payload *payload)
+void arch_prog_run(struct prog *prog)
 {
-	stage_exit(prog_entry(&payload->prog));
+	stage_exit(prog_entry(prog));
+}
+
+void arch_payload_run(struct payload *payload)
+{
+	arch_prog_run(&payload->prog);
 }
diff --git a/src/arch/riscv/Makefile.inc b/src/arch/riscv/Makefile.inc
index 242fd404e6..30c84a93e2 100644
--- a/src/arch/riscv/Makefile.inc
+++ b/src/arch/riscv/Makefile.inc
@@ -31,6 +31,7 @@ ifeq ($(CONFIG_ARCH_BOOTBLOCK_RISCV),y)
 bootblock_lds = $(src)/arch/riscv/bootblock.ld
 
 bootblock-y = bootblock.S stages.c
+bootblock-y += boot.c
 bootblock-y += rom_media.c
 bootblock-y += \
 	$(top)/src/lib/memchr.c \
@@ -52,6 +53,7 @@ endif
 ################################################################################
 ifeq ($(CONFIG_ARCH_ROMSTAGE_RISCV),y)
 
+romstage-y += boot.c
 romstage-y += stages.c
 romstage-y += rom_media.c
 romstage-y += \
diff --git a/src/arch/riscv/boot.c b/src/arch/riscv/boot.c
index 4c2be031c2..f66c02ff2e 100644
--- a/src/arch/riscv/boot.c
+++ b/src/arch/riscv/boot.c
@@ -17,11 +17,17 @@
  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
-#include <arch/stages.h>
 #include <program_loading.h>
 
-void arch_payload_run(const struct payload *payload)
+void arch_prog_run(struct prog *prog)
 {
-//	uart_rx_byte(0);
-	stage_exit(prog_entry(&payload->prog));
+	void (*doit)(void *);
+
+	doit = prog_entry(prog);
+	doit(prog_entry_arg(prog));
+}
+
+void arch_payload_run(struct payload *payload)
+{
+	arch_prog_run(&payload->prog);
 }
diff --git a/src/arch/x86/boot/Makefile.inc b/src/arch/x86/boot/Makefile.inc
index 9f06b4f4eb..7810ba8368 100644
--- a/src/arch/x86/boot/Makefile.inc
+++ b/src/arch/x86/boot/Makefile.inc
@@ -2,6 +2,7 @@
 ifeq ($(CONFIG_ARCH_ROMSTAGE_X86_32),y)
 
 romstage-y += cbmem.c
+romstage-y += boot.c
 
 endif # CONFIG_ARCH_ROMSTAGE_X86_32
 
diff --git a/src/arch/x86/boot/boot.c b/src/arch/x86/boot/boot.c
index 1fb36e2827..d5365842d2 100644
--- a/src/arch/x86/boot/boot.c
+++ b/src/arch/x86/boot/boot.c
@@ -123,7 +123,7 @@ static void jmp_payload(void *entry, unsigned long buffer, unsigned long size)
 		);
 }
 
-void arch_payload_run(const struct payload *payload)
+void arch_payload_run(struct payload *payload)
 {
 	if (IS_ENABLED(CONFIG_RELOCATABLE_RAMSTAGE))
 		jmp_payload_no_bounce_buffer(prog_entry(&payload->prog));
@@ -132,3 +132,11 @@ void arch_payload_run(const struct payload *payload)
 				(uintptr_t)payload->bounce.data,
 				payload->bounce.size);
 }
+
+void arch_prog_run(struct prog *prog)
+{
+	__asm__ volatile (
+		"jmp  *%%edi\n"
+		:: "D"(prog_entry(prog))
+	);
+}