diff --git a/arch/arm/mach-rockchip/make_fit_atf.py b/arch/arm/mach-rockchip/make_fit_atf.py
index d1faff1957..99cf6f43dc 100755
--- a/arch/arm/mach-rockchip/make_fit_atf.py
+++ b/arch/arm/mach-rockchip/make_fit_atf.py
@@ -194,8 +194,11 @@ def get_bl31_segments_info(bl31_file_name):
 
 def main():
     uboot_elf="./u-boot"
-    bl31_elf="./bl31.elf"
     FIT_ITS=sys.stdout
+    if "BL31" in os.environ:
+        bl31_elf=os.getenv("BL31");
+    else:
+        sys.exit("ERROR: Please export BL31 file, check doc/README.rockchip")
 
     opts, args = getopt.getopt(sys.argv[1:], "o:u:b:h")
     for opt, val in opts:
diff --git a/doc/README.rockchip b/doc/README.rockchip
index e8f6a2dfba..0db4771162 100644
--- a/doc/README.rockchip
+++ b/doc/README.rockchip
@@ -138,8 +138,8 @@ For example:
      => make realclean
      => make CROSS_COMPILE=aarch64-linux-gnu- PLAT=rk3399
 
-     (copy bl31.elf U-Boot root dir)
-     => cp build/rk3399/release/bl31/bl31.elf /path/to/u-boot
+     (export bl31.elf)
+     => export BL31=/path/to/arm-trusted-firmware/build/rk3399/release/bl31/bl31.elf
 
    - Compile U-Boot