spl: stash bootstage info before jump to next stage

Since we may jump to next stage like ATF/OP-TEE instead of U-Boot,
we need to stash the bootstage info before it.

Signed-off-by: Kever Yang <kever.yang@rock-chips.com>
Reviewed-by: Simon Glass <sjg@chromium.org>
diff --git a/common/spl/spl.c b/common/spl/spl.c
index ce9819e..4afbe97 100644
--- a/common/spl/spl.c
+++ b/common/spl/spl.c
@@ -424,13 +424,6 @@
 	debug("SPL malloc() used %#lx bytes (%ld KB)\n", gd->malloc_ptr,
 	      gd->malloc_ptr / 1024);
 #endif
-
-	if (CONFIG_IS_ENABLED(ATF_SUPPORT)) {
-		debug("loaded - jumping to U-Boot via ATF BL31.\n");
-		bl31_entry();
-	}
-
-	debug("loaded - jumping to U-Boot...\n");
 #ifdef CONFIG_BOOTSTAGE_STASH
 	int ret;
 
@@ -440,6 +433,13 @@
 	if (ret)
 		debug("Failed to stash bootstage: err=%d\n", ret);
 #endif
+
+	if (CONFIG_IS_ENABLED(ATF_SUPPORT)) {
+		debug("loaded - jumping to U-Boot via ATF BL31.\n");
+		bl31_entry();
+	}
+
+	debug("loaded - jumping to U-Boot...\n");
 	spl_board_prepare_for_boot();
 	jump_to_image_no_args(&spl_image);
 }