--- a/board/mediatek/mt7981/mt7981_rfb.c
+++ b/board/mediatek/mt7981/mt7981_rfb.c
@@ -10,7 +10,9 @@
 #include <env.h>
 #include <init.h>
 #include <asm/global_data.h>
+#include <asm/io.h>
 #include <linux/delay.h>
+#include <linux/libfdt.h>
 
 #ifndef CONFIG_RESET_BUTTON_LABEL
 #define CONFIG_RESET_BUTTON_LABEL "reset"
@@ -80,3 +82,54 @@ int board_nmbm_init(void)
 
 	return 0;
 }
+
+#define	MT7981_BOOT_NOR		0
+#define	MT7981_BOOT_SPIM_NAND	1 /* ToDo: fallback to SD */
+#define	MT7981_BOOT_EMMC	2
+#define	MT7981_BOOT_SNFI_NAND	3 /* ToDo (treated as SD) */
+
+int ft_system_setup(void *blob, struct bd_info *bd)
+{
+	const u32 *media_handle_p;
+	int chosen, len, ret;
+	const char *media;
+	u32 media_handle;
+
+	switch ((readl(0x11d006f0) & 0xc0) >> 6) {
+	case MT7981_BOOT_NOR:
+		media = "rootdisk-nor";
+		break
+		;;
+	case MT7981_BOOT_SPIM_NAND:
+		media = "rootdisk-spim-nand";
+		break
+		;;
+	case MT7981_BOOT_EMMC:
+		media = "rootdisk-emmc";
+		break
+		;;
+	case MT7981_BOOT_SNFI_NAND:
+		media = "rootdisk-sd";
+		break
+		;;
+	}
+
+	chosen = fdt_path_offset(blob, "/chosen");
+	if (chosen <= 0)
+		return 0;
+
+	media_handle_p = fdt_getprop(blob, chosen, media, &len);
+	if (media_handle_p <= 0 || len != 4)
+		return 0;
+
+	media_handle = *media_handle_p;
+	ret = fdt_setprop(blob, chosen, "rootdisk", &media_handle, sizeof(media_handle));
+	if (ret) {
+		printf("cannot set media phandle %s as rootdisk /chosen node\n", media);
+		return ret;
+	}
+
+	printf("set /chosen/rootdisk to bootrom media: %s (phandle 0x%08x)\n", media, fdt32_to_cpu(media_handle));
+
+	return 0;
+}
--- a/configs/mt7981_emmc_rfb_defconfig
+++ b/configs/mt7981_emmc_rfb_defconfig
@@ -114,3 +114,4 @@ CONFIG_FAT_WRITE=y
 CONFIG_HEXDUMP=y
 # CONFIG_EFI_LOADER is not set
 CONFIG_LMB_MAX_REGIONS=64
+CONFIG_OF_SYSTEM_SETUP=y
--- a/configs/mt7981_nor_rfb_defconfig
+++ b/configs/mt7981_nor_rfb_defconfig
@@ -126,3 +126,4 @@ CONFIG_DM_SPI=y
 CONFIG_MTK_SPIM=y
 CONFIG_HEXDUMP=y
 CONFIG_LMB_MAX_REGIONS=64
+CONFIG_OF_SYSTEM_SETUP=y
--- a/configs/mt7981_rfb_defconfig
+++ b/configs/mt7981_rfb_defconfig
@@ -135,3 +135,4 @@ CONFIG_MTK_SPIM=y
 CONFIG_HEXDUMP=y
 # CONFIG_EFI_LOADER is not set
 CONFIG_LMB_MAX_REGIONS=64
+CONFIG_OF_SYSTEM_SETUP=y
--- a/configs/mt7981_sd_rfb_defconfig
+++ b/configs/mt7981_sd_rfb_defconfig
@@ -114,3 +114,4 @@ CONFIG_FAT_WRITE=y
 CONFIG_HEXDUMP=y
 # CONFIG_EFI_LOADER is not set
 CONFIG_LMB_MAX_REGIONS=64
+CONFIG_OF_SYSTEM_SETUP=y
--- a/configs/mt7981_snfi_nand_rfb_defconfig
+++ b/configs/mt7981_snfi_nand_rfb_defconfig
@@ -119,3 +119,4 @@ CONFIG_DM_SERIAL=y
 CONFIG_MTK_SERIAL=y
 CONFIG_HEXDUMP=y
 CONFIG_LMB_MAX_REGIONS=64
+CONFIG_OF_SYSTEM_SETUP=y
