diff --git a/arch/arm/dts/sunxi-u-boot.dtsi b/arch/arm/dts/sunxi-u-boot.dtsi
index b7195c5..b8a0d4b 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_MACH_SUN50I_H6
+#define BL31_ADDR 0x104000
+#else
+#define BL31_ADDR  0x44000
+#endif
+
 / {
 	aliases {
 		mmc1 = &mmc2;
@@ -42,14 +48,8 @@
 					type = "firmware";
 					arch = "arm64";
 					compression = "none";
-/* TODO: Do this with an overwrite in this board's dtb? */
-#ifdef CONFIG_MACH_SUN50I_H6
-					load = <0x104000>;
-					entry = <0x104000>;
-#else
-					load = <0x44000>;
-					entry = <0x44000>;
-#endif
+					load = <BL31_ADDR>;
+					entry = <BL31_ADDR>;
 
 					atf-bl31 {
 						filename = "bl31.bin";