diff --git a/arch/arm/dts/sunxi-u-boot.dtsi b/arch/arm/dts/sunxi-u-boot.dtsi
index 7a8764e..0909a67 100644
--- a/arch/arm/dts/sunxi-u-boot.dtsi
+++ b/arch/arm/dts/sunxi-u-boot.dtsi
@@ -1,5 +1,11 @@
 #include <config.h>
 
+#ifdef CONFIG_ARM64
+#define ARCH "arm64"
+#else
+#define ARCH "arm"
+#endif
+
 / {
 	aliases {
 		mmc0 = &mmc0;
@@ -34,30 +40,34 @@
 			filename = "spl/sunxi-spl.bin";
 		};
 
-#ifdef CONFIG_ARM64
+#ifdef CONFIG_SPL_LOAD_FIT
 		fit {
-			description = "Configuration to load ATF before U-Boot";
+			description = "Configuration to load U-Boot and firmware";
 			#address-cells = <1>;
 			fit,fdt-list = "of-list";
 
 			images {
 				uboot {
-					description = "U-Boot (64-bit)";
+					description = "U-Boot";
 					type = "standalone";
 					os = "u-boot";
-					arch = "arm64";
+					arch = ARCH;
 					compression = "none";
 					load = <CONFIG_TEXT_BASE>;
+#if CONFIG_SUNXI_BL31_BASE == 0
+					entry = <CONFIG_TEXT_BASE>;
+#endif
 
 					u-boot-nodtb {
 					};
 				};
 
+#if CONFIG_SUNXI_BL31_BASE
 				atf {
 					description = "ARM Trusted Firmware";
 					type = "firmware";
 					os = "arm-trusted-firmware";
-					arch = "arm64";
+					arch = ARCH;
 					compression = "none";
 					load = <CONFIG_SUNXI_BL31_BASE>;
 					entry = <CONFIG_SUNXI_BL31_BASE>;
@@ -67,6 +77,7 @@
 						missing-msg = "atf-bl31-sunxi";
 					};
 				};
+#endif
 
 #if CONFIG_SUNXI_SCP_BASE
 				scp {
@@ -95,7 +106,11 @@
 
 				@config-SEQ {
 					description = "NAME";
+#if CONFIG_SUNXI_BL31_BASE
 					firmware = "atf";
+#else
+					firmware = "uboot";
+#endif
 #if CONFIG_SUNXI_SCP_BASE
 					loadables = "scp", "uboot";
 #else