Merge patch series "lmb: use a single API for all allocations"

Sughosh Ganu <sughosh.ganu@linaro.org> says:

The LMB module has a bunch for API's which are used for allocating
memory. There are a couple of API's for requesting memory, and two
more for reserving regions of memory. Replace these different API's
with a single one, lmb_alloc_mem(). The type of allocation to be made
is specified through one of the parameters to the function.

Additionally, the two API's for reserving regions of memory,
lmb_reserve() and lmb_alloc_addr() are the same with one
difference. One can reserve any memory region with lmb_reserve(),
while lmb_alloc_addr() actually checks that the memory region being
requested is part of the LMB memory map. Reserving memory that is not
part of the LMB memory map is pretty futile -- the allocation
functions do not allocate memory which has not been added to the LMB
memory map.

This series also removes the functionality allowing for reserving
memory regions outside the LMB memory map. Any request for reserving a
region of memory outside the LMB memory map now returns an -EINVAL
error.

Certain places in the common code using the LMB API's were not
checking the return value of the functions. Checks have been added for
them. There are some calls being made from the architecture/platform
specific code which too do not check the return value. Those have been
kept the same, as I do not have the platform with me to check if it
causes any issues on those platforms.

In addition, there is a patch which refactors code in
lmb_overlaps_region() and lmb_can_reserve_region() so that both
functionalities can be put in a single function, lmb_overlap_checks().

Finally, a new patch has been added which checks the return value of
the lmb allocation function before copying the device-tree to the
allocated address.

Link: https://lore.kernel.org/r/20250617104346.1379981-1-sughosh.ganu@linaro.org
[trini: Rework arch/arm/mach-snapdragon/board.c merge]
Signed-off-by: Tom Rini <trini@konsulko.com>

20 files changed, 580 insertions, 282 deletions

diff --git a/arch/arm/mach-apple/board.c b/arch/arm/mach-apple/board.c
index 2604c5a710e..4cd8979bdc2 100644
--- a/arch/arm/mach-apple/board.c
+++ b/arch/arm/mach-apple/board.c
@@ -773,22 +773,31 @@ u64 get_page_table_size(void)
 
 #define KERNEL_COMP_SIZE	SZ_128M
 
+#define lmb_alloc(size, addr) lmb_alloc_mem(LMB_MEM_ALLOC_ANY, SZ_2M, addr, size, LMB_NONE)
+
 int board_late_init(void)
 {
 	u32 status = 0;
+	phys_addr_t addr;
 
 	/* somewhat based on the Linux Kernel boot requirements:
 	 * align by 2M and maximal FDT size 2M
 	 */
-	status |= env_set_hex("loadaddr", lmb_alloc(SZ_1G, SZ_2M));
-	status |= env_set_hex("fdt_addr_r", lmb_alloc(SZ_2M, SZ_2M));
-	status |= env_set_hex("kernel_addr_r", lmb_alloc(SZ_128M, SZ_2M));
-	status |= env_set_hex("ramdisk_addr_r", lmb_alloc(SZ_1G, SZ_2M));
-	status |= env_set_hex("kernel_comp_addr_r",
-			      lmb_alloc(KERNEL_COMP_SIZE, SZ_2M));
-	status |= env_set_hex("kernel_comp_size", KERNEL_COMP_SIZE);
-	status |= env_set_hex("scriptaddr", lmb_alloc(SZ_4M, SZ_2M));
-	status |= env_set_hex("pxefile_addr_r", lmb_alloc(SZ_4M, SZ_2M));
+	status |= !lmb_alloc(SZ_1G, &addr) ? env_set_hex("loadaddr", addr) : 1;
+	status |= !lmb_alloc(SZ_2M, &addr) ?
+		env_set_hex("fdt_addr_r", addr) : 1;
+	status |= !lmb_alloc(SZ_128M, &addr) ?
+		env_set_hex("kernel_addr_r", addr) : 1;
+	status |= !lmb_alloc(SZ_1G, &addr) ?
+		env_set_hex("ramdisk_addr_r", addr) : 1;
+	status |= !lmb_alloc(KERNEL_COMP_SIZE, &addr) ?
+		env_set_hex("kernel_comp_addr_r", addr) : 1;
+	status |= !lmb_alloc(KERNEL_COMP_SIZE, &addr) ?
+		env_set_hex("kernel_comp_size", addr) : 1;
+	status |= !lmb_alloc(SZ_4M, &addr) ?
+		env_set_hex("scriptaddr", addr) : 1;
+	status |= !lmb_alloc(SZ_4M, &addr) ?
+		env_set_hex("pxefile_addr_r", addr) : 1;
 
 	if (status)
 		log_warning("late_init: Failed to set run time variables\n");
diff --git a/arch/arm/mach-mediatek/tzcfg.c b/arch/arm/mach-mediatek/tzcfg.c
index 71982ba4d20..c8fe8ac0e9b 100644
--- a/arch/arm/mach-mediatek/tzcfg.c
+++ b/arch/arm/mach-mediatek/tzcfg.c
@@ -173,6 +173,7 @@ phys_addr_t board_get_usable_ram_top(phys_size_t total_size)
 
 int arch_misc_init(void)
 {
+	phys_addr_t addr;
 	struct arm_smccc_res res;
 
 	/*
@@ -180,11 +181,14 @@ int arch_misc_init(void)
 	 * there's no need to check the result
 	 */
 	arm_smccc_smc(MTK_SIP_GET_BL31_REGION, 0, 0, 0, 0, 0, 0, 0, &res);
-	lmb_reserve(res.a1, res.a2, LMB_NOMAP);
+	addr = (phys_addr_t)res.a1;
+	lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr, res.a2, LMB_NOMAP);
 
 	arm_smccc_smc(MTK_SIP_GET_BL32_REGION, 0, 0, 0, 0, 0, 0, 0, &res);
+	addr = (phys_addr_t)res.a1;
 	if (!res.a0 && res.a1 && res.a2)
-		lmb_reserve(res.a1, res.a2, LMB_NOMAP);
+		lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr, res.a2,
+			      LMB_NOMAP);
 
 #if IS_ENABLED(CONFIG_CMD_PSTORE)
 	char cmd[64];
diff --git a/arch/arm/mach-snapdragon/board.c b/arch/arm/mach-snapdragon/board.c
index 8f7fc19cadc..ec51ebbbe7f 100644
--- a/arch/arm/mach-snapdragon/board.c
+++ b/arch/arm/mach-snapdragon/board.c
@@ -517,38 +517,51 @@ void __weak qcom_late_init(void)
 #define FASTBOOT_BUF_SIZE 0
 #endif
 
-#define addr_alloc(size) lmb_alloc(size, SZ_2M)
+#define lmb_alloc(size, addr) lmb_alloc_mem(LMB_MEM_ALLOC_ANY, SZ_2M, addr, size, LMB_NONE)
 
 /* Stolen from arch/arm/mach-apple/board.c */
 int board_late_init(void)
 {
-	u32 status = 0;
+	u32 status = 0, fdt_status = 0;
 	phys_addr_t addr;
 	struct fdt_header *fdt_blob = (struct fdt_header *)gd->fdt_blob;
 
 	/* We need to be fairly conservative here as we support boards with just 1G of TOTAL RAM */
-	addr = addr_alloc(SZ_128M);
+	status |= !lmb_alloc(SZ_128M, &addr) ?
+		env_set_hex("loadaddr", addr) : 1;
 	status |= env_set_hex("kernel_addr_r", addr);
-	status |= env_set_hex("loadaddr", addr);
-	status |= env_set_hex("ramdisk_addr_r", addr_alloc(SZ_128M));
-	status |= env_set_hex("kernel_comp_addr_r", addr_alloc(KERNEL_COMP_SIZE));
-	status |= env_set_hex("kernel_comp_size", KERNEL_COMP_SIZE);
+	status |= !lmb_alloc(SZ_128M, &addr) ?
+		env_set_hex("ramdisk_addr_r", addr) : 1;
+	status |= !lmb_alloc(KERNEL_COMP_SIZE, &addr) ?
+		env_set_hex("kernel_comp_addr_r", addr) : 1;
+	status |= !lmb_alloc(KERNEL_COMP_SIZE, &addr) ?
+		env_set_hex("kernel_comp_size", addr) : 1;
+	status |= !lmb_alloc(SZ_4M, &addr) ?
+		env_set_hex("scriptaddr", addr) : 1;
+	status |= !lmb_alloc(SZ_4M, &addr) ?
+		env_set_hex("pxefile_addr_r", addr) : 1;
+
 	if (IS_ENABLED(CONFIG_FASTBOOT)) {
-		addr = addr_alloc(FASTBOOT_BUF_SIZE);
-		status |= env_set_hex("fastboot_addr_r", addr);
+		status |= !lmb_alloc(FASTBOOT_BUF_SIZE, &addr) ?
+			env_set_hex("fastboot_addr_r", addr) : 1;
 		/* override loadaddr for memory rich soc */
-		status |= env_set_hex("loadaddr", addr);
+		status |= !lmb_alloc(SZ_128M, &addr) ?
+			env_set_hex("loadaddr", addr) : 1;
 	}
-	status |= env_set_hex("scriptaddr", addr_alloc(SZ_4M));
-	status |= env_set_hex("pxefile_addr_r", addr_alloc(SZ_4M));
-	addr = addr_alloc(SZ_2M);
-	status |= env_set_hex("fdt_addr_r", addr);
 
-	if (status)
+	fdt_status |= !lmb_alloc(SZ_2M, &addr) ?
+		env_set_hex("fdt_addr_r", addr) : 1;
+
+	if (status || fdt_status)
 		log_warning("%s: Failed to set run time variables\n", __func__);
 
 	/* By default copy U-Boots FDT, it will be used as a fallback */
-	memcpy((void *)addr, (void *)gd->fdt_blob, fdt32_to_cpu(fdt_blob->totalsize));
+	if (fdt_status)
+		log_warning("%s: Failed to reserve memory for copying FDT\n",
+			    __func__);
+	else
+		memcpy((void *)addr, (void *)gd->fdt_blob,
+		       fdt32_to_cpu(fdt_blob->totalsize));
 
 	configure_env();
 	qcom_late_init();
diff --git a/arch/powerpc/cpu/mpc85xx/mp.c b/arch/powerpc/cpu/mpc85xx/mp.c
index 8918a401fac..bee6758dc9a 100644
--- a/arch/powerpc/cpu/mpc85xx/mp.c
+++ b/arch/powerpc/cpu/mpc85xx/mp.c
@@ -410,9 +410,9 @@ static void plat_mp_up(unsigned long bootpg, unsigned int pagesize)
 
 void cpu_mp_lmb_reserve(void)
 {
-	u32 bootpg = determine_mp_bootpg(NULL);
+	phys_addr_t bootpg = determine_mp_bootpg(NULL);
 
-	lmb_reserve(bootpg, 4096, LMB_NONE);
+	lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &bootpg, 4096, LMB_NONE);
 }
 
 void setup_mp(void)
diff --git a/arch/powerpc/lib/misc.c b/arch/powerpc/lib/misc.c
index 7e303419624..fc10ae50cf8 100644
--- a/arch/powerpc/lib/misc.c
+++ b/arch/powerpc/lib/misc.c
@@ -36,11 +36,12 @@ int arch_misc_init(void)
 		size = min(size, (ulong)CFG_SYS_LINUX_LOWMEM_MAX_SIZE);
 
 		if (size < bootm_size) {
-			ulong base = bootmap_base + size;
+			phys_addr_t base = bootmap_base + size;
 
 			printf("WARNING: adjusting available memory from 0x%lx to 0x%llx\n",
 			       size, (unsigned long long)bootm_size);
-			lmb_reserve(base, bootm_size - size, LMB_NONE);
+			lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &base,
+				      bootm_size - size, LMB_NONE);
 		}
 
 #ifdef CONFIG_MP
diff --git a/boot/bootm.c b/boot/bootm.c
index 108ca7fb472..4bdca22ea8c 100644
--- a/boot/bootm.c
+++ b/boot/bootm.c
@@ -623,12 +623,16 @@ static int bootm_load_os(struct bootm_headers *images, int boot_progress)
 	 */
 	if (os.type == IH_TYPE_KERNEL_NOLOAD && os.comp != IH_COMP_NONE) {
 		ulong req_size = ALIGN(image_len * 4, SZ_1M);
+		phys_addr_t addr;
 
-		load = lmb_alloc(req_size, SZ_2M);
-		if (!load)
+		err = lmb_alloc_mem(LMB_MEM_ALLOC_ANY, SZ_2M, &addr,
+				    req_size, LMB_NONE);
+		if (err)
 			return 1;
-		os.load = load;
-		images->ep = load;
+
+		load = (ulong)addr;
+		os.load = (ulong)addr;
+		images->ep = (ulong)addr;
 		debug("Allocated %lx bytes at %lx for kernel (size %lx) decompression\n",
 		      req_size, load, image_len);
 	}
@@ -698,9 +702,18 @@ static int bootm_load_os(struct bootm_headers *images, int boot_progress)
 		images->os.end = relocated_addr + image_size;
 	}
 
-	if (CONFIG_IS_ENABLED(LMB))
-		lmb_reserve(images->os.load, (load_end - images->os.load),
-			    LMB_NONE);
+	if (CONFIG_IS_ENABLED(LMB)) {
+		phys_addr_t load;
+
+		load = (phys_addr_t)images->os.load;
+		err = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &load,
+				    (load_end - images->os.load), LMB_NONE);
+		if (err) {
+			log_err("Unable to allocate memory %#lx for loading OS\n",
+				images->os.load);
+			return 1;
+		}
+	}
 
 	return 0;
 }
diff --git a/boot/image-board.c b/boot/image-board.c
index 514f8e63f9c..005d60caf5c 100644
--- a/boot/image-board.c
+++ b/boot/image-board.c
@@ -16,6 +16,7 @@
 #include <fpga.h>
 #include <image.h>
 #include <init.h>
+#include <lmb.h>
 #include <log.h>
 #include <mapmem.h>
 #include <rtc.h>
@@ -538,6 +539,7 @@ int boot_get_ramdisk(char const *select, struct bootm_headers *images,
 int boot_ramdisk_high(ulong rd_data, ulong rd_len, ulong *initrd_start,
 		      ulong *initrd_end)
 {
+	int err;
 	char	*s;
 	phys_addr_t initrd_high;
 	int	initrd_copy_to_ram = 1;
@@ -559,25 +561,30 @@ int boot_ramdisk_high(ulong rd_data, ulong rd_len, ulong *initrd_start,
 
 	if (rd_data) {
 		if (!initrd_copy_to_ram) {	/* zero-copy ramdisk support */
+			phys_addr_t initrd_addr;
+
 			debug("   in-place initrd\n");
 			*initrd_start = rd_data;
 			*initrd_end = rd_data + rd_len;
-			lmb_reserve(rd_data, rd_len, LMB_NONE);
+			initrd_addr = (phys_addr_t)rd_data;
+			err = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &initrd_addr,
+					    rd_len, LMB_NONE);
+			if (err) {
+				puts("in-place initrd alloc failed\n");
+				goto error;
+			}
 		} else {
-			if (initrd_high)
-				*initrd_start =
-					(ulong)lmb_alloc_base(rd_len,
-								    0x1000,
-								    initrd_high,
-								    LMB_NONE);
-			else
-				*initrd_start = (ulong)lmb_alloc(rd_len,
-								 0x1000);
+			enum lmb_mem_type type = initrd_high ?
+				LMB_MEM_ALLOC_MAX : LMB_MEM_ALLOC_ANY;
 
-			if (*initrd_start == 0) {
+			err = lmb_alloc_mem(type, 0x1000, &initrd_high, rd_len,
+					    LMB_NONE);
+			if (err) {
 				puts("ramdisk - allocation error\n");
 				goto error;
 			}
+
+			*initrd_start = (ulong)initrd_high;
 			bootstage_mark(BOOTSTAGE_ID_COPY_RAMDISK);
 
 			*initrd_end = *initrd_start + rd_len;
@@ -828,9 +835,10 @@ int boot_get_loadable(struct bootm_headers *images)
  */
 int boot_get_cmdline(ulong *cmd_start, ulong *cmd_end)
 {
-	int barg;
+	int barg, err;
 	char *cmdline;
 	char *s;
+	phys_addr_t addr;
 
 	/*
 	 * Help the compiler detect that this function is only called when
@@ -840,12 +848,14 @@ int boot_get_cmdline(ulong *cmd_start, ulong *cmd_end)
 		return 0;
 
 	barg = IF_ENABLED_INT(CONFIG_SYS_BOOT_GET_CMDLINE, CONFIG_SYS_BARGSIZE);
-	cmdline = (char *)(ulong)lmb_alloc_base(barg, 0xf,
-				env_get_bootm_mapsize() + env_get_bootm_low(),
-				LMB_NONE);
-	if (!cmdline)
+	addr = env_get_bootm_mapsize() + env_get_bootm_low();
+
+	err = lmb_alloc_mem(LMB_MEM_ALLOC_MAX, 0xf, &addr, barg, LMB_NONE);
+	if (err)
 		return -1;
 
+	cmdline = (char *)(uintptr_t)addr;
+
 	s = env_get("bootargs");
 	if (!s)
 		s = "";
@@ -874,14 +884,16 @@ int boot_get_cmdline(ulong *cmd_start, ulong *cmd_end)
  */
 int boot_get_kbd(struct bd_info **kbd)
 {
-	*kbd = (struct bd_info *)(ulong)lmb_alloc_base(sizeof(struct bd_info),
-							     0xf,
-							     env_get_bootm_mapsize() +
-							     env_get_bootm_low(),
-							     LMB_NONE);
-	if (!*kbd)
+	int err;
+	phys_addr_t addr;
+
+	addr = env_get_bootm_mapsize() + env_get_bootm_low();
+	err = lmb_alloc_mem(LMB_MEM_ALLOC_MAX, 0xf, &addr,
+			    sizeof(struct bd_info), LMB_NONE);
+	if (err)
 		return -1;
 
+	*kbd = (struct bd_info *)(uintptr_t)addr;
 	**kbd = *gd->bd;
 
 	debug("## kernel board info at 0x%08lx\n", (ulong)*kbd);
diff --git a/boot/image-fdt.c b/boot/image-fdt.c
index 8f718ad29f6..97b6385ab7c 100644
--- a/boot/image-fdt.c
+++ b/boot/image-fdt.c
@@ -72,13 +72,15 @@ static const struct legacy_img_hdr *image_get_fdt(ulong fdt_addr)
 static void boot_fdt_reserve_region(u64 addr, u64 size, u32 flags)
 {
 	long ret;
+	phys_addr_t rsv_addr;
 
-	ret = lmb_reserve(addr, size, flags);
+	rsv_addr = (phys_addr_t)addr;
+	ret = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &rsv_addr, size, flags);
 	if (!ret) {
 		debug("   reserving fdt memory region: addr=%llx size=%llx flags=%x\n",
 		      (unsigned long long)addr,
 		      (unsigned long long)size, flags);
-	} else if (ret != -EEXIST) {
+	} else if (ret != -EEXIST && ret != -EINVAL) {
 		puts("ERROR: reserving fdt memory region failed ");
 		printf("(addr=%llx size=%llx flags=%x)\n",
 		       (unsigned long long)addr,
@@ -155,7 +157,7 @@ void boot_fdt_add_mem_rsv_regions(void *fdt_blob)
  */
 int boot_relocate_fdt(char **of_flat_tree, ulong *of_size)
 {
-	u64	start, size, usable, addr, low, mapsize;
+	u64	start, size, usable, low, mapsize;
 	void	*fdt_blob = *of_flat_tree;
 	void	*of_start = NULL;
 	char	*fdt_high;
@@ -163,6 +165,7 @@ int boot_relocate_fdt(char **of_flat_tree, ulong *of_size)
 	int	bank;
 	int	err;
 	int	disable_relocation = 0;
+	phys_addr_t addr;
 
 	/* nothing to do */
 	if (*of_size == 0)
@@ -180,23 +183,32 @@ int boot_relocate_fdt(char **of_flat_tree, ulong *of_size)
 	/* If fdt_high is set use it to select the relocation address */
 	fdt_high = env_get("fdt_high");
 	if (fdt_high) {
-		ulong desired_addr = hextoul(fdt_high, NULL);
+		ulong high_addr = hextoul(fdt_high, NULL);
 
-		if (desired_addr == ~0UL) {
+		if (high_addr == ~0UL) {
 			/* All ones means use fdt in place */
 			of_start = fdt_blob;
-			lmb_reserve(map_to_sysmem(of_start), of_len, LMB_NONE);
-			disable_relocation = 1;
-		} else if (desired_addr) {
-			addr = lmb_alloc_base(of_len, 0x1000, desired_addr,
-					      LMB_NONE);
-			of_start = map_sysmem(addr, of_len);
-			if (of_start == NULL) {
-				puts("Failed using fdt_high value for Device Tree");
+			addr = map_to_sysmem(fdt_blob);
+			err = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr,
+					    of_len, LMB_NONE);
+			if (err) {
+				printf("Failed to reserve memory for fdt at %#llx\n",
+				       (u64)addr);
 				goto error;
 			}
+
+			disable_relocation = 1;
 		} else {
-			addr = lmb_alloc(of_len, 0x1000);
+			enum lmb_mem_type type = high_addr ?
+				LMB_MEM_ALLOC_MAX : LMB_MEM_ALLOC_ANY;
+
+			addr = high_addr;
+			err = lmb_alloc_mem(type, 0x1000, &addr, of_len,
+					    LMB_NONE);
+			if (err) {
+				puts("Failed to allocate memory for Device Tree relocation\n");
+				goto error;
+			}
 			of_start = map_sysmem(addr, of_len);
 		}
 	} else {
@@ -218,11 +230,15 @@ int boot_relocate_fdt(char **of_flat_tree, ulong *of_size)
 			 * for LMB allocation.
 			 */
 			usable = min(start + size, low + mapsize);
-			addr = lmb_alloc_base(of_len, 0x1000, usable, LMB_NONE);
-			of_start = map_sysmem(addr, of_len);
-			/* Allocation succeeded, use this block. */
-			if (of_start != NULL)
-				break;
+			addr = usable;
+			err = lmb_alloc_mem(LMB_MEM_ALLOC_MAX, 0x1000,
+					    &addr, of_len, LMB_NONE);
+			if (!err) {
+				of_start = map_sysmem(addr, of_len);
+				/* Allocation succeeded, use this block. */
+				if (of_start)
+					break;
+			}
 
 			/*
 			 * Reduce the mapping size in the next bank
@@ -674,7 +690,7 @@ int image_setup_libfdt(struct bootm_headers *images, void *blob, bool lmb)
 
 	/* Delete the old LMB reservation */
 	if (CONFIG_IS_ENABLED(LMB) && lmb)
-		lmb_free(map_to_sysmem(blob), fdt_totalsize(blob));
+		lmb_free(map_to_sysmem(blob), fdt_totalsize(blob), LMB_NONE);
 
 	ret = fdt_shrink_to_minimum(blob, 0);
 	if (ret < 0)
@@ -682,8 +698,17 @@ int image_setup_libfdt(struct bootm_headers *images, void *blob, bool lmb)
 	of_size = ret;
 
 	/* Create a new LMB reservation */
-	if (CONFIG_IS_ENABLED(LMB) && lmb)
-		lmb_reserve(map_to_sysmem(blob), of_size, LMB_NONE);
+	if (CONFIG_IS_ENABLED(LMB) && lmb) {
+		phys_addr_t fdt_addr;
+
+		fdt_addr = map_to_sysmem(blob);
+		ret = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &fdt_addr,
+				    of_size, LMB_NONE);
+		if (ret) {
+			printf("Failed to reserve memory for the fdt at %#llx\n",
+			       (u64)fdt_addr);
+		}
+	}
 
 #if defined(CONFIG_ARCH_KEYSTONE)
 	if (IS_ENABLED(CONFIG_OF_BOARD_SETUP))
diff --git a/cmd/booti.c b/cmd/booti.c
index 7e6d9426299..e6f67d6e136 100644
--- a/cmd/booti.c
+++ b/cmd/booti.c
@@ -30,6 +30,7 @@ static int booti_start(struct bootm_info *bmi)
 	uint8_t *temp;
 	ulong dest;
 	ulong dest_end;
+	phys_addr_t ep_addr;
 	unsigned long comp_len;
 	unsigned long decomp_len;
 	int ctype;
@@ -88,7 +89,14 @@ static int booti_start(struct bootm_info *bmi)
 	images->os.start = relocated_addr;
 	images->os.end = relocated_addr + image_size;
 
-	lmb_reserve(images->ep, le32_to_cpu(image_size), LMB_NONE);
+	ep_addr = (phys_addr_t)images->ep;
+	ret = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &ep_addr,
+			    le32_to_cpu(image_size), LMB_NONE);
+	if (ret) {
+		printf("Failed to allocate memory for the image at %#llx\n",
+		       (unsigned long long)images->ep);
+		return 1;
+	}
 
 	/*
 	 * Handle the BOOTM_STATE_FINDOTHER state ourselves as we do not
diff --git a/cmd/bootz.c b/cmd/bootz.c
index 99318ff213f..44af7d012aa 100644
--- a/cmd/bootz.c
+++ b/cmd/bootz.c
@@ -28,6 +28,7 @@ static int bootz_start(struct cmd_tbl *cmdtp, int flag, int argc,
 {
 	ulong zi_start, zi_end;
 	struct bootm_info bmi;
+	phys_addr_t ep_addr;
 	int ret;
 
 	bootm_init(&bmi);
@@ -56,7 +57,14 @@ static int bootz_start(struct cmd_tbl *cmdtp, int flag, int argc,
 	if (ret != 0)
 		return 1;
 
-	lmb_reserve(images->ep, zi_end - zi_start, LMB_NONE);
+	ep_addr = (phys_addr_t)images->ep;
+	ret = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &ep_addr, zi_end - zi_start,
+			    LMB_NONE);
+	if (ret) {
+		printf("Failed to allocate memory for the image at %#llx\n",
+		       (unsigned long long)images->ep);
+		return 1;
+	}
 
 	/*
 	 * Handle the BOOTM_STATE_FINDOTHER state ourselves as we do not
diff --git a/cmd/load.c b/cmd/load.c
index 899bb4f598e..159767aa7f7 100644
--- a/cmd/load.c
+++ b/cmd/load.c
@@ -178,17 +178,20 @@ static ulong load_serial(long offset)
 #endif
 		    {
 			void *dst;
+			phys_addr_t dst_addr;
 
-			ret = lmb_reserve(store_addr, binlen, LMB_NONE);
+			dst_addr = (phys_addr_t)store_addr;
+			ret = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &dst_addr,
+					    binlen, LMB_NONE);
 			if (ret) {
 				printf("\nCannot overwrite reserved area (%08lx..%08lx)\n",
 					store_addr, store_addr + binlen);
 				return ret;
 			}
-			dst = map_sysmem(store_addr, binlen);
+			dst = map_sysmem(dst_addr, binlen);
 			memcpy(dst, binbuf, binlen);
 			unmap_sysmem(dst);
-			lmb_free(store_addr, binlen);
+			lmb_free(dst_addr, binlen, LMB_NONE);
 		    }
 		    if ((store_addr) < start_addr)
 			start_addr = store_addr;
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 506843ed74a..cf9d21e4c1c 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -17,7 +17,6 @@ U-Boot API documentation
    interrupt
    led
    linker_lists
-   lmb
    logging
    nvmem
    part
diff --git a/doc/api/lmb.rst b/doc/api/lmb.rst
deleted file mode 100644
index 2095bfa1618..00000000000
--- a/doc/api/lmb.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-.. SPDX-License-Identifier: GPL-2.0+
-
-Logical memory blocks
-=====================
-
-.. kernel-doc:: include/lmb.h
-   :internal:
diff --git a/doc/develop/index.rst b/doc/develop/index.rst
index 0c83ef109ab..3c044e67927 100644
--- a/doc/develop/index.rst
+++ b/doc/develop/index.rst
@@ -46,6 +46,7 @@ Implementation
    cedit
    event
    global_data
+   lmb
    logging
    makefiles
    menus
diff --git a/doc/develop/lmb.rst b/doc/develop/lmb.rst
new file mode 100644
index 00000000000..b9d0f09c2bb
--- /dev/null
+++ b/doc/develop/lmb.rst
@@ -0,0 +1,166 @@
+.. SPDX-License-Identifier: GPL-2.0+
+
+Logical Memory Blocks (LMB)
+===========================
+
+U-Boot has support for reserving chunks of memory which is primarily
+used for loading images to the DRAM memory, before these are booted,
+or written to non-volatile storage medium. This functionality is
+provided through the Logical Memory Blocks (LMB) module.
+
+Introduction
+------------
+
+The LMB module manages allocation requests for memory region not
+occupied by the U-Boot image. Allocation requests that are made
+through malloc() and similar functions result in memory getting
+allocated from the heap region, which is part of the U-Boot
+image. Typically, the heap memory is a few MiB in size. Loading an
+image like the linux kernel might require lot more memory than what
+the heap can provide. Such allocations are usually handled through the
+LMB module.
+
+The U-Boot image typically gets relocated to the top of the usable
+DRAM memory region. A typical memory layout looks as follows::
+
+
+
+
+
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+       ---            +--------------+  <--- U-Boot ram top
+        |             |              |
+        |             |    Text      |
+        |             +--------------+
+        |             |              |
+        |             |    Data      |
+        |             +--------------+
+        |             |              |
+        |             |    BSS       |
+  U-Boot Image        +--------------+
+        |             |              |
+        |             |    Heap      |
+        |             |              |
+        |             +--------------+
+        |             |              |
+        |             |              |
+        |             |    Stack     |
+        |             |              |
+        |             |              |
+       ---            +--------------+
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      |              |
+                      +--------------+  <--- ram start
+
+
+
+The region of memory below the U-Boot image is the one controlled by
+the LMB module.
+
+
+Types of LMB Allocations
+------------------------
+
+There are two classes of allocation requests that get made to the LMB
+module. One type of allocation requests are requesting memory of a
+particular number of bytes. This type of allocation is similar to that
+done using the malloc type of function calls. The other type of
+allocations, are requests made for a specific memory address. The
+second type of allocations are usually made for loading images to a
+particular memory address.
+
+
+LMB design Pre 2025.01
+----------------------
+
+The earlier versions of U-Boot (pre 2025.01 release)
+had a local memory map based LMB implementation whereby it was
+possible to declare the LMB map inside a function or a C file. This
+design resulted in temporary, non-global LMB maps, which also allowed
+for re-use of memory. This meant that it was possible to use a region
+of memory to load some image, and subsequently the same region of
+memory could be used for loading a different image. A typical example
+of this usage would be loading an image to a memory address, followed
+by writing that image to some non-volatile storage medium. Once this
+is done, the same address can be used for loading a different image
+and then writing it to it's non-volatile storage
+destination. Typically, environment variables like `loadaddr`,
+`kernel_addr_r`, `ramdisk_addr_r` are used for loading images to
+memory regions.
+
+
+Current LMB implementation
+--------------------------
+
+Changes were made in the 2025.01 release to make the LMB memory map
+global and persistent. With this, the LMB memory map is the same
+across all of U-Boot, and also persists as long as U-Boot is
+active. Even with this change, there has been consistency as far as
+re-use of memory is concerned to maintain backward compatibility. It
+is allowed for re-requesting the same region of memory if the memory
+region has a particular attribute (LMB_NONE).
+
+As part of the platform boot, DRAM memory available for use in U-Boot
+gets added to the LMB memory map. Any allocation requests made
+subsequently will be made from this memory added as part of the board
+init.
+
+
+Allocation API
+--------------
+
+Any request for non-heap memory can be made through the LMB allocation
+API.
+
+.. code-block:: c
+
+	int lmb_alloc_mem(enum lmb_mem_type type, u64 align,
+			  phys_addr_t *addr, phys_size_t size,
+			  u32 flags);
+
+Correspondingly, the allocated memory can be free'd
+
+.. code-block:: c
+
+	long lmb_free(phys_addr_t base, phys_size_t size, u32 flags);
+
+For a detailed API description, please refer to the header file.
+
+
+UEFI allocations with LMB as the backend
+----------------------------------------
+
+The UEFI specification describes boot-time API's for allocation of
+memory. These API's use the same memory that is being used by the LMB
+module. Pre 2025.01 release, there wasn't any synchronisation between
+the EFI sub-system and the LMB module about the memory that was
+getting allocated by each of these modules. This was the primary
+reason for making the LMB memory map global and persistent. With this
+change, the EFI memory allocation API's have also been changed to use
+the LMB module as the backend for the allocation requests. Any other
+sub-system which might wish to use the same memory region for it's use
+can then use the LMB as the backend for the memory allocations and
+it's associated book-keeping.
+
+
+API documentation
+-----------------
+
+.. kernel-doc:: include/lmb.h
+
diff --git a/fs/fs.c b/fs/fs.c
index 1f36872fb9a..2650328b358 100644
--- a/fs/fs.c
+++ b/fs/fs.c
@@ -580,6 +580,7 @@ static int fs_read_lmb_check(const char *filename, ulong addr, loff_t offset,
 	int ret;
 	loff_t size;
 	loff_t read_len;
+	phys_addr_t read_addr;
 
 	/* get the actual size of the file */
 	ret = info->size(filename, &size);
@@ -597,7 +598,9 @@ static int fs_read_lmb_check(const char *filename, ulong addr, loff_t offset,
 
 	lmb_dump_all();
 
-	if (!lmb_alloc_addr(addr, read_len, LMB_NONE))
+	read_addr = (phys_addr_t)addr;
+	if (!lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &read_addr, read_len,
+			   LMB_NONE))
 		return 0;
 
 	log_err("** Reading file would overwrite reserved memory **\n");
diff --git a/include/lmb.h b/include/lmb.h
index 606a92cca48..5d5f037ccb9 100644
--- a/include/lmb.h
+++ b/include/lmb.h
@@ -32,6 +32,18 @@
 #define LMB_NONOTIFY BIT(3)
 
 /**
+ * enum lmb_mem_type - type of memory allocation request
+ * @LMB_MEM_ALLOC_ADDR:	request for a particular region of memory
+ * @LMB_MEM_ALLOC_ANY:	allocate any available memory region
+ * @LMB_MEM_ALLOC_MAX:	allocate memory below a particular address
+ */
+enum lmb_mem_type {
+	LMB_MEM_ALLOC_ADDR = 1,
+	LMB_MEM_ALLOC_ANY,
+	LMB_MEM_ALLOC_MAX,
+};
+
+/**
  * enum lmb_map_op - memory map operation
  */
 enum lmb_map_op {
@@ -68,6 +80,37 @@ struct lmb {
 };
 
 /**
+ * lmb_alloc_mem() - Request LMB memory
+ * @type:		Type of memory allocation request
+ * @align:		Alignment of the memory region requested(0 for none)
+ * @addr:		Base address of the allocated memory region
+ * @size:		Size in bytes of the allocation request
+ * @flags:		Memory region attributes to be set
+ *
+ * Allocate a region of memory where the allocation is based on the parameters
+ * that have been passed to the function.The first parameter specifies the
+ * type of allocation that is being requested. The second parameter, @align
+ * is used to specify if the allocation is to be made with a particular
+ * alignment. Use 0 for no alignment requirements.
+ *
+ * The allocated address is returned through the @addr parameter when @type
+ * is @LMB_MEM_ALLOC_ANY or @LMB_MEM_ALLOC_MAX. If @type is
+ * @LMB_MEM_ALLOC_ADDR the @addr parameter would contain the address being
+ * requested.
+ *
+ * The flags parameter is used to specify the memory attributes of the
+ * requested region.
+ *
+ * Return: 0 on success, -ve value on failure
+ *
+ * When the allocation is of type @LMB_MEM_ALLOC_ADDR, the return value can
+ * be -EINVAL if the requested memory region is not part of the LMB memory
+ * map, and -EEXIST if the requested region is already allocated.
+ */
+int lmb_alloc_mem(enum lmb_mem_type type, u64 align, phys_addr_t *addr,
+		  phys_size_t size, u32 flags);
+
+/**
  * lmb_init() - Initialise the LMB module.
  *
  * Return: 0 on success, negative error code on failure.
@@ -81,65 +124,11 @@ struct lmb {
  */
 int lmb_init(void);
 
-/**
- * lmb_add_memory() - Add memory range for LMB allocations.
- *
- * Add the entire available memory range to the pool of memory that
- * can be used by the LMB module for allocations.
- */
-void lmb_add_memory(void);
-
 long lmb_add(phys_addr_t base, phys_size_t size);
 
-/**
- * lmb_reserve() - Reserve one region with a specific flags bitfield
- * @base: Base address of the memory region
- * @size: Size of the memory region
- * @flags: Flags for the memory region
- *
- * Return:
- * * %0		- Added successfully, or it's already added (only if LMB_NONE)
- * * %-EEXIST	- The region is already added, and flags != LMB_NONE
- * * %-1	- Failure
- */
-long lmb_reserve(phys_addr_t base, phys_size_t size, u32 flags);
-
-phys_addr_t lmb_alloc(phys_size_t size, ulong align);
 phys_size_t lmb_get_free_size(phys_addr_t addr);
 
 /**
- * lmb_alloc_base() - Allocate specified memory region with specified
- *			    attributes
- * @size: Size of the region requested
- * @align: Alignment of the memory region requested
- * @max_addr: Maximum address of the requested region
- * @flags: Memory region attributes to be set
- *
- * Allocate a region of memory with the attributes specified through the
- * parameter. The max_addr parameter is used to specify the maximum address
- * below which the requested region should be allocated.
- *
- * Return: Base address on success, 0 on error.
- */
-phys_addr_t lmb_alloc_base(phys_size_t size, ulong align, phys_addr_t max_addr,
-			   uint flags);
-
-/**
- * lmb_alloc_addr() - Allocate specified memory address with specified attributes
- *
- * @base: Base Address requested
- * @size: Size of the region requested
- * @flags: Memory region attributes to be set
- *
- * Allocate a region of memory with the attributes specified through the
- * parameter. The base parameter is used to specify the base address
- * of the requested region.
- *
- * Return: 0 on success -1 on error
- */
-int lmb_alloc_addr(phys_addr_t base, phys_size_t size, u32 flags);
-
-/**
  * lmb_is_reserved_flags() - Test if address is in reserved region with flag
  *			     bits set
  * @addr: Address to be tested
@@ -153,16 +142,14 @@ int lmb_alloc_addr(phys_addr_t base, phys_size_t size, u32 flags);
 int lmb_is_reserved_flags(phys_addr_t addr, int flags);
 
 /**
- * lmb_free_flags() - Free up a region of memory
+ * lmb_free() - Free up a region of memory
  * @base: Base Address of region to be freed
  * @size: Size of the region to be freed
  * @flags: Memory region attributes
  *
  * Return: 0 on success, negative error code on failure.
  */
-long lmb_free_flags(phys_addr_t base, phys_size_t size, uint flags);
-
-long lmb_free(phys_addr_t base, phys_size_t size);
+long lmb_free(phys_addr_t base, phys_size_t size, u32 flags);
 
 void lmb_dump_all(void);
 void lmb_dump_all_force(void);
@@ -175,7 +162,7 @@ void lmb_pop(struct lmb *store);
 
 static inline int lmb_read_check(phys_addr_t addr, phys_size_t len)
 {
-	return lmb_alloc_addr(addr, len, LMB_NONE);
+	return lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr, len, LMB_NONE);
 }
 
 /**
diff --git a/lib/efi_loader/efi_memory.c b/lib/efi_loader/efi_memory.c
index 0abb1f6159a..0828a47da61 100644
--- a/lib/efi_loader/efi_memory.c
+++ b/lib/efi_loader/efi_memory.c
@@ -454,6 +454,7 @@ efi_status_t efi_allocate_pages(enum efi_allocate_type type,
 				enum efi_memory_type memory_type,
 				efi_uintn_t pages, uint64_t *memory)
 {
+	int err;
 	u64 efi_addr, len;
 	uint flags;
 	efi_status_t ret;
@@ -475,17 +476,18 @@ efi_status_t efi_allocate_pages(enum efi_allocate_type type,
 	switch (type) {
 	case EFI_ALLOCATE_ANY_PAGES:
 		/* Any page */
-		addr = (u64)lmb_alloc_base(len, EFI_PAGE_SIZE,
-						 LMB_ALLOC_ANYWHERE, flags);
-		if (!addr)
+		err = lmb_alloc_mem(LMB_MEM_ALLOC_ANY, EFI_PAGE_SIZE, &addr,
+				    len, flags);
+		if (err)
 			return EFI_OUT_OF_RESOURCES;
 		break;
 	case EFI_ALLOCATE_MAX_ADDRESS:
 		/* Max address */
 		addr = map_to_sysmem((void *)(uintptr_t)*memory);
-		addr = (u64)lmb_alloc_base(len, EFI_PAGE_SIZE, addr,
-						 flags);
-		if (!addr)
+
+		err = lmb_alloc_mem(LMB_MEM_ALLOC_MAX, EFI_PAGE_SIZE, &addr,
+				    len, flags);
+		if (err)
 			return EFI_OUT_OF_RESOURCES;
 		break;
 	case EFI_ALLOCATE_ADDRESS:
@@ -493,7 +495,7 @@ efi_status_t efi_allocate_pages(enum efi_allocate_type type,
 			return EFI_NOT_FOUND;
 
 		addr = map_to_sysmem((void *)(uintptr_t)*memory);
-		if (lmb_alloc_addr(addr, len, flags))
+		if (lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr, len, flags))
 			return EFI_NOT_FOUND;
 		break;
 	default:
@@ -506,7 +508,7 @@ efi_status_t efi_allocate_pages(enum efi_allocate_type type,
 	ret = efi_update_memory_map(efi_addr, pages, memory_type, true, false);
 	if (ret != EFI_SUCCESS) {
 		/* Map would overlap, bail out */
-		lmb_free_flags(addr, (u64)pages << EFI_PAGE_SHIFT, flags);
+		lmb_free(addr, (u64)pages << EFI_PAGE_SHIFT, flags);
 		unmap_sysmem((void *)(uintptr_t)efi_addr);
 		return  EFI_OUT_OF_RESOURCES;
 	}
@@ -546,8 +548,8 @@ efi_status_t efi_free_pages(uint64_t memory, efi_uintn_t pages)
 	 * been mapped with map_sysmem() from efi_allocate_pages(). Convert
 	 * it back to an address LMB understands
 	 */
-	status = lmb_free_flags(map_to_sysmem((void *)(uintptr_t)memory), len,
-				LMB_NOOVERWRITE);
+	status = lmb_free(map_to_sysmem((void *)(uintptr_t)memory), len,
+			  LMB_NOOVERWRITE);
 	if (status)
 		return EFI_NOT_FOUND;
 
diff --git a/lib/lmb.c b/lib/lmb.c
index bb6f232f6bc..45b26512a5b 100644
--- a/lib/lmb.c
+++ b/lib/lmb.c
@@ -317,8 +317,34 @@ static long _lmb_free(struct alist *lmb_rgn_lst, phys_addr_t base,
 				    rgn[i].flags);
 }
 
-static long lmb_overlaps_region(struct alist *lmb_rgn_lst, phys_addr_t base,
-				phys_size_t size)
+/**
+ * lmb_overlap_checks() - perform checks to see if region can be allocated or reserved
+ * @lmb_rgn_lst: list of LMB regions
+ * @base: base address of region to be checked
+ * @size: size of region to be checked
+ * @flags: flag of the region to be checked (only for reservation requests)
+ * @alloc: if checks are to be done for allocation or reservation request
+ *
+ * Check if the region passed to the function overlaps with any one of
+ * the regions of the passed lmb region list.
+ *
+ * If the @alloc flag is set to true, this check stops as soon an
+ * overlapping region is found. The function can also be called to
+ * check if a reservation request can be satisfied, by setting
+ * @alloc to false. In that case, the function then iterates through
+ * all the regions in the list passed to ensure that the requested
+ * region does not overlap with any existing regions. An overlap is
+ * allowed only when the flag of the requested region and the existing
+ * region is LMB_NONE.
+ *
+ * Return: index of the overlapping region, -1 if no overlap is found
+ *
+ * When the function is called for a reservation request check, -1 will
+ * also be returned when there is an allowed overlap, i.e. requested
+ * region and existing regions have flags as LMB_NONE.
+ */
+static long lmb_overlap_checks(struct alist *lmb_rgn_lst, phys_addr_t base,
+			       phys_size_t size, u32 flags, bool alloc)
 {
 	unsigned long i;
 	struct lmb_region *rgn = lmb_rgn_lst->data;
@@ -326,9 +352,12 @@ static long lmb_overlaps_region(struct alist *lmb_rgn_lst, phys_addr_t base,
 	for (i = 0; i < lmb_rgn_lst->count; i++) {
 		phys_addr_t rgnbase = rgn[i].base;
 		phys_size_t rgnsize = rgn[i].size;
+		u32 rgnflags = rgn[i].flags;
 
-		if (lmb_addrs_overlap(base, size, rgnbase, rgnsize))
-			break;
+		if (lmb_addrs_overlap(base, size, rgnbase, rgnsize)) {
+			if (alloc || flags != LMB_NONE || flags != rgnflags)
+				break;
+		}
 	}
 
 	return (i < lmb_rgn_lst->count) ? i : -1;
@@ -390,7 +419,8 @@ phys_addr_t io_lmb_alloc(struct lmb *io_lmb, phys_size_t size, ulong align)
 		base = ALIGN_DOWN(lmbbase + lmbsize - size, align);
 
 		while (base && lmbbase <= base) {
-			rgn = lmb_overlaps_region(&io_lmb->used_mem, base, size);
+			rgn = lmb_overlap_checks(&io_lmb->used_mem, base, size,
+						 LMB_NOOVERWRITE, true);
 			if (rgn < 0) {
 				/* This area isn't reserved, take it */
 				if (lmb_add_region_flags(&io_lmb->used_mem, base,
@@ -488,6 +518,21 @@ void lmb_dump_all(void)
 #endif
 }
 
+static long lmb_reserve(phys_addr_t base, phys_size_t size, u32 flags)
+{
+	long ret = 0;
+	struct alist *lmb_rgn_lst = &lmb.used_mem;
+
+	if (lmb_overlap_checks(lmb_rgn_lst, base, size, flags, false) != -1)
+		return -EEXIST;
+
+	ret = lmb_add_region_flags(lmb_rgn_lst, base, size, flags);
+	if (ret)
+		return ret;
+
+	return lmb_map_update_notify(base, size, LMB_MAP_OP_RESERVE, flags);
+}
+
 static void lmb_reserve_uboot_region(void)
 {
 	int bank;
@@ -557,40 +602,7 @@ static __maybe_unused void lmb_reserve_common_spl(void)
 	}
 }
 
-/**
- * lmb_can_reserve_region() - check if the region can be reserved
- * @base: base address of region to be reserved
- * @size: size of region to be reserved
- * @flags: flag of the region to be reserved
- *
- * Go through all the reserved regions and ensure that the requested
- * region does not overlap with any existing regions. An overlap is
- * allowed only when the flag of the request region and the existing
- * region is LMB_NONE.
- *
- * Return: true if region can be reserved, false otherwise
- */
-static bool lmb_can_reserve_region(phys_addr_t base, phys_size_t size,
-				   u32 flags)
-{
-	uint i;
-	struct lmb_region *lmb_reserved = lmb.used_mem.data;
-
-	for (i = 0; i < lmb.used_mem.count; i++) {
-		u32 rgnflags = lmb_reserved[i].flags;
-		phys_addr_t rgnbase = lmb_reserved[i].base;
-		phys_size_t rgnsize = lmb_reserved[i].size;
-
-		if (lmb_addrs_overlap(base, size, rgnbase, rgnsize)) {
-			if (flags != LMB_NONE || flags != rgnflags)
-				return false;
-		}
-	}
-
-	return true;
-}
-
-void lmb_add_memory(void)
+static void lmb_add_memory(void)
 {
 	int i;
 	phys_addr_t bank_end;
@@ -640,8 +652,7 @@ long lmb_add(phys_addr_t base, phys_size_t size)
 	return lmb_map_update_notify(base, size, LMB_MAP_OP_ADD, LMB_NONE);
 }
 
-long lmb_free_flags(phys_addr_t base, phys_size_t size,
-		    uint flags)
+long lmb_free(phys_addr_t base, phys_size_t size, u32 flags)
 {
 	long ret;
 
@@ -652,36 +663,18 @@ long lmb_free_flags(phys_addr_t base, phys_size_t size,
 	return lmb_map_update_notify(base, size, LMB_MAP_OP_FREE, flags);
 }
 
-long lmb_free(phys_addr_t base, phys_size_t size)
-{
-	return lmb_free_flags(base, size, LMB_NONE);
-}
-
-long lmb_reserve(phys_addr_t base, phys_size_t size, u32 flags)
-{
-	long ret = 0;
-	struct alist *lmb_rgn_lst = &lmb.used_mem;
-
-	if (!lmb_can_reserve_region(base, size, flags))
-		return -EEXIST;
-
-	ret = lmb_add_region_flags(lmb_rgn_lst, base, size, flags);
-	if (ret)
-		return ret;
-
-	return lmb_map_update_notify(base, size, LMB_MAP_OP_RESERVE, flags);
-}
-
-static phys_addr_t _lmb_alloc_base(phys_size_t size, ulong align,
-				   phys_addr_t max_addr, u32 flags)
+static int _lmb_alloc_base(phys_size_t size, ulong align,
+			   phys_addr_t *addr, u32 flags)
 {
 	int ret;
 	long i, rgn;
+	phys_addr_t max_addr;
 	phys_addr_t base = 0;
 	phys_addr_t res_base;
 	struct lmb_region *lmb_used = lmb.used_mem.data;
 	struct lmb_region *lmb_memory = lmb.available_mem.data;
 
+	max_addr = *addr;
 	for (i = lmb.available_mem.count - 1; i >= 0; i--) {
 		phys_addr_t lmbbase = lmb_memory[i].base;
 		phys_size_t lmbsize = lmb_memory[i].size;
@@ -702,7 +695,8 @@ static phys_addr_t _lmb_alloc_base(phys_size_t size, ulong align,
 		}
 
 		while (base && lmbbase <= base) {
-			rgn = lmb_overlaps_region(&lmb.used_mem, base, size);
+			rgn = lmb_overlap_checks(&lmb.used_mem, base, size,
+						 LMB_NOOVERWRITE, true);
 			if (rgn < 0) {
 				/* This area isn't reserved, take it */
 				if (lmb_add_region_flags(&lmb.used_mem, base,
@@ -714,8 +708,8 @@ static phys_addr_t _lmb_alloc_base(phys_size_t size, ulong align,
 							    flags);
 				if (ret)
 					return ret;
-
-				return base;
+				*addr = base;
+				return 0;
 			}
 
 			res_base = lmb_used[rgn].base;
@@ -728,27 +722,17 @@ static phys_addr_t _lmb_alloc_base(phys_size_t size, ulong align,
 	log_debug("%s: Failed to allocate 0x%lx bytes below 0x%lx\n",
 		  __func__, (ulong)size, (ulong)max_addr);
 
-	return 0;
-}
-
-phys_addr_t lmb_alloc(phys_size_t size, ulong align)
-{
-	return _lmb_alloc_base(size, align, LMB_ALLOC_ANYWHERE, LMB_NONE);
-}
-
-phys_addr_t lmb_alloc_base(phys_size_t size, ulong align, phys_addr_t max_addr,
-			   uint flags)
-{
-	return _lmb_alloc_base(size, align, max_addr, flags);
+	return -1;
 }
 
-int lmb_alloc_addr(phys_addr_t base, phys_size_t size, u32 flags)
+static int _lmb_alloc_addr(phys_addr_t base, phys_size_t size, u32 flags)
 {
 	long rgn;
 	struct lmb_region *lmb_memory = lmb.available_mem.data;
 
 	/* Check if the requested address is in one of the memory regions */
-	rgn = lmb_overlaps_region(&lmb.available_mem, base, size);
+	rgn = lmb_overlap_checks(&lmb.available_mem, base, size,
+				 LMB_NOOVERWRITE, true);
 	if (rgn >= 0) {
 		/*
 		 * Check if the requested end address is in the same memory
@@ -756,14 +740,40 @@ int lmb_alloc_addr(phys_addr_t base, phys_size_t size, u32 flags)
 		 */
 		if (lmb_addrs_overlap(lmb_memory[rgn].base,
 				      lmb_memory[rgn].size,
-				      base + size - 1, 1)) {
+				      base + size - 1, 1))
 			/* ok, reserve the memory */
-			if (!lmb_reserve(base, size, flags))
-				return 0;
-		}
+			return lmb_reserve(base, size, flags);
 	}
 
-	return -1;
+	return -EINVAL;
+}
+
+int lmb_alloc_mem(enum lmb_mem_type type, u64 align, phys_addr_t *addr,
+		  phys_size_t size, u32 flags)
+{
+	int ret = -1;
+
+	if (!size)
+		return 0;
+
+	if (!addr)
+		return -EINVAL;
+
+	switch (type) {
+	case LMB_MEM_ALLOC_ANY:
+		*addr = LMB_ALLOC_ANYWHERE;
+	case LMB_MEM_ALLOC_MAX:
+		ret = _lmb_alloc_base(size, align, addr, flags);
+		break;
+	case LMB_MEM_ALLOC_ADDR:
+		ret = _lmb_alloc_addr(*addr, size, flags);
+		break;
+	default:
+		log_debug("%s: Invalid memory allocation type requested %d\n",
+			  __func__, type);
+	}
+
+	return ret;
 }
 
 /* Return number of bytes from a given address that are free */
@@ -775,7 +785,8 @@ phys_size_t lmb_get_free_size(phys_addr_t addr)
 	struct lmb_region *lmb_memory = lmb.available_mem.data;
 
 	/* check if the requested address is in the memory regions */
-	rgn = lmb_overlaps_region(&lmb.available_mem, addr, 1);
+	rgn = lmb_overlap_checks(&lmb.available_mem, addr, 1, LMB_NOOVERWRITE,
+				 true);
 	if (rgn >= 0) {
 		for (i = 0; i < lmb.used_mem.count; i++) {
 			if (addr < lmb_used[i].base) {
diff --git a/test/lib/lmb.c b/test/lib/lmb.c
index 3bf558f7f4f..b6259bef442 100644
--- a/test/lib/lmb.c
+++ b/test/lib/lmb.c
@@ -71,6 +71,45 @@ static int setup_lmb_test(struct unit_test_state *uts, struct lmb *store,
 	return 0;
 }
 
+static int lmb_reserve(phys_addr_t addr, phys_size_t size, u32 flags)
+{
+	int err;
+
+	err = lmb_alloc_mem(LMB_MEM_ALLOC_ADDR, 0, &addr, size, flags);
+	if (err)
+		return err;
+
+	return 0;
+}
+
+static phys_addr_t lmb_alloc(phys_size_t size, ulong align)
+{
+	int err;
+	phys_addr_t addr;
+
+	err = lmb_alloc_mem(LMB_MEM_ALLOC_ANY, align, &addr, size, LMB_NONE);
+	if (err)
+		return 0;
+
+	return addr;
+}
+
+static phys_addr_t lmb_alloc_base(phys_size_t size, ulong align,
+				  phys_addr_t max_addr, u32 flags)
+{
+	int err;
+	phys_addr_t addr;
+
+	addr = max_addr;
+	err = lmb_alloc_mem(LMB_MEM_ALLOC_MAX, align, &addr, size, flags);
+	if (err)
+		return 0;
+
+	return addr;
+}
+
+#define lmb_alloc_addr(addr, size, flags) lmb_reserve(addr, size, flags)
+
 static int test_multi_alloc(struct unit_test_state *uts, const phys_addr_t ram,
 			    const phys_size_t ram_size, const phys_addr_t ram0,
 			    const phys_size_t ram0_size,
@@ -143,7 +182,7 @@ static int test_multi_alloc(struct unit_test_state *uts, const phys_addr_t ram,
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 0x10000 + 8, ram_end - 8, 8, 0, 0);
 
-	ret = lmb_free(a, 4);
+	ret = lmb_free(a, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 0x10000 + 8, ram_end - 8, 4, 0, 0);
@@ -152,12 +191,12 @@ static int test_multi_alloc(struct unit_test_state *uts, const phys_addr_t ram,
 	ut_asserteq(a, a2);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 0x10000 + 8, ram_end - 8, 8, 0, 0);
-	ret = lmb_free(a2, 4);
+	ret = lmb_free(a2, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 0x10000 + 8, ram_end - 8, 4, 0, 0);
 
-	ret = lmb_free(b, 4);
+	ret = lmb_free(b, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 3,
 		   alloc_64k_addr - 8, 4, alloc_64k_addr, 0x10000,
@@ -167,17 +206,17 @@ static int test_multi_alloc(struct unit_test_state *uts, const phys_addr_t ram,
 	ut_asserteq(b, b2);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 0x10000 + 8, ram_end - 8, 4, 0, 0);
-	ret = lmb_free(b2, 4);
+	ret = lmb_free(b2, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 3,
 		   alloc_64k_addr - 8, 4, alloc_64k_addr, 0x10000,
 		   ram_end - 8, 4);
 
-	ret = lmb_free(c, 4);
+	ret = lmb_free(c, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 2,
 		   alloc_64k_addr - 8, 4, alloc_64k_addr, 0x10000, 0, 0);
-	ret = lmb_free(d, 4);
+	ret = lmb_free(d, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, 0, 0, 1, alloc_64k_addr, 0x10000,
 		   0, 0, 0, 0);
@@ -281,7 +320,7 @@ static int test_bigblock(struct unit_test_state *uts, const phys_addr_t ram)
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, a,
 		   big_block_size + 0x10000, 0, 0, 0, 0);
 
-	ret = lmb_free(a, big_block_size);
+	ret = lmb_free(a, big_block_size, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, alloc_64k_addr, 0x10000,
 		   0, 0, 0, 0);
@@ -353,12 +392,12 @@ static int test_noreserved(struct unit_test_state *uts, const phys_addr_t ram,
 			   - alloc_size_aligned, alloc_size, 0, 0);
 	}
 	/* and free them */
-	ret = lmb_free(b, alloc_size);
+	ret = lmb_free(b, alloc_size, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1,
 		   ram + ram_size - alloc_size_aligned,
 		   alloc_size, 0, 0, 0, 0);
-	ret = lmb_free(a, alloc_size);
+	ret = lmb_free(a, alloc_size, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 0, 0, 0, 0, 0, 0, 0);
 
@@ -369,7 +408,7 @@ static int test_noreserved(struct unit_test_state *uts, const phys_addr_t ram,
 		   ram + ram_size - alloc_size_aligned,
 		   alloc_size, 0, 0, 0, 0);
 	/* and free it */
-	ret = lmb_free(b, alloc_size);
+	ret = lmb_free(b, alloc_size, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 0, 0, 0, 0, 0, 0, 0);
 
@@ -437,12 +476,12 @@ static int lib_test_lmb_at_0(struct unit_test_state *uts)
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, a, ram_size - 4,
 		   0, 0, 0, 0);
 	/* check that this was an error by freeing b */
-	ret = lmb_free(b, 4);
+	ret = lmb_free(b, 4, LMB_NONE);
 	ut_asserteq(ret, -1);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, a, ram_size - 4,
 		   0, 0, 0, 0);
 
-	ret = lmb_free(a, ram_size - 4);
+	ret = lmb_free(a, ram_size - 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 0, 0, 0, 0, 0, 0, 0);
 
@@ -568,20 +607,20 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NONE);
 	ut_asserteq(b, 0);
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NOOVERWRITE);
-	ut_asserteq(b, -1);
+	ut_asserteq(b, -EEXIST);
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NONE);
 	ut_asserteq(b, 0);
 	b = lmb_alloc_addr(alloc_addr_a, 0x2000, LMB_NONE);
 	ut_asserteq(b, 0);
-	ret = lmb_free(alloc_addr_a, 0x2000);
+	ret = lmb_free(alloc_addr_a, 0x2000, LMB_NONE);
 	ut_asserteq(ret, 0);
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NOOVERWRITE);
 	ut_asserteq(b, 0);
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NONE);
-	ut_asserteq(b, -1);
+	ut_asserteq(b, -EEXIST);
 	b = lmb_alloc_addr(alloc_addr_a, 0x1000, LMB_NOOVERWRITE);
-	ut_asserteq(b, -1);
-	ret = lmb_free(alloc_addr_a, 0x1000);
+	ut_asserteq(b, -EEXIST);
+	ret = lmb_free(alloc_addr_a, 0x1000, LMB_NONE);
 	ut_asserteq(ret, 0);
 
 	/*
@@ -599,13 +638,13 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 		   alloc_addr_a + 0x4000, 0x1000, 0, 0);
 
 	c = lmb_alloc_addr(alloc_addr_a + 0x1000, 0x5000, LMB_NONE);
-	ut_asserteq(c, -1);
+	ut_asserteq(c, -EEXIST);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 2, alloc_addr_a, 0x1000,
 		   alloc_addr_a + 0x4000, 0x1000, 0, 0);
 
-	ret = lmb_free(alloc_addr_a, 0x1000);
+	ret = lmb_free(alloc_addr_a, 0x1000, LMB_NONE);
 	ut_asserteq(ret, 0);
-	ret = lmb_free(alloc_addr_a + 0x4000, 0x1000);
+	ret = lmb_free(alloc_addr_a + 0x4000, 0x1000, LMB_NOOVERWRITE);
 	ut_asserteq(ret, 0);
 
 	/*
@@ -628,7 +667,7 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, alloc_addr_a, 0x6000,
 		   0, 0, 0, 0);
 
-	ret = lmb_free(alloc_addr_a, 0x6000);
+	ret = lmb_free(alloc_addr_a, 0x6000, LMB_NONE);
 	ut_asserteq(ret, 0);
 
 	/*
@@ -646,13 +685,13 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 		   alloc_addr_a + 0x4000, 0x1000, 0, 0);
 
 	c = lmb_alloc_addr(alloc_addr_a + 0x1000, 0x5000, LMB_NOOVERWRITE);
-	ut_asserteq(c, -1);
+	ut_asserteq(c, -EEXIST);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 2, alloc_addr_a, 0x1000,
 		   alloc_addr_a + 0x4000, 0x1000, 0, 0);
 
-	ret = lmb_free(alloc_addr_a, 0x1000);
+	ret = lmb_free(alloc_addr_a, 0x1000, LMB_NOOVERWRITE);
 	ut_asserteq(ret, 0);
-	ret = lmb_free(alloc_addr_a + 0x4000, 0x1000);
+	ret = lmb_free(alloc_addr_a + 0x4000, 0x1000, LMB_NOOVERWRITE);
 	ut_asserteq(ret, 0);
 
 	/*  reserve 3 blocks */
@@ -693,7 +732,8 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 		   0, 0, 0, 0);
 
 	/* free thge allocation from d */
-	ret = lmb_free(alloc_addr_c + 0x10000, ram_end - alloc_addr_c - 0x10000);
+	ret = lmb_free(alloc_addr_c + 0x10000, ram_end - alloc_addr_c - 0x10000,
+		       LMB_NONE);
 	ut_asserteq(ret, 0);
 
 	/* allocate at 3 points in free range */
@@ -702,7 +742,7 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	ut_asserteq(d, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 2, ram, 0x18010000,
 		   ram_end - 4, 4, 0, 0);
-	ret = lmb_free(ram_end - 4, 4);
+	ret = lmb_free(ram_end - 4, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, ram, 0x18010000,
 		   0, 0, 0, 0);
@@ -711,7 +751,7 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	ut_asserteq(d, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 2, ram, 0x18010000,
 		   ram_end - 128, 4, 0, 0);
-	ret = lmb_free(ram_end - 128, 4);
+	ret = lmb_free(ram_end - 128, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, ram, 0x18010000,
 		   0, 0, 0, 0);
@@ -720,13 +760,13 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	ut_asserteq(d, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, ram, 0x18010004,
 		   0, 0, 0, 0);
-	ret = lmb_free(alloc_addr_c + 0x10000, 4);
+	ret = lmb_free(alloc_addr_c + 0x10000, 4, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, ram, 0x18010000,
 		   0, 0, 0, 0);
 
 	/* allocate at the bottom a was assigned to ram at the top */
-	ret = lmb_free(ram, alloc_addr_a - ram);
+	ret = lmb_free(ram, alloc_addr_a - ram, LMB_NONE);
 	ut_asserteq(ret, 0);
 	ASSERT_LMB(mem_lst, used_lst, ram, ram_size, 1, ram + 0x8000000,
 		   0x10010000, 0, 0, 0, 0);
@@ -739,11 +779,11 @@ static int test_alloc_addr(struct unit_test_state *uts, const phys_addr_t ram)
 	/* check that allocating outside memory fails */
 	if (ram_end != 0) {
 		ret = lmb_alloc_addr(ram_end, 1, LMB_NONE);
-		ut_asserteq(ret, -1);
+		ut_asserteq(ret, -EINVAL);
 	}
 	if (ram != 0) {
 		ret = lmb_alloc_addr(ram - 1, 1, LMB_NONE);
-		ut_asserteq(ret, -1);
+		ut_asserteq(ret, -EINVAL);
 	}
 
 	lmb_pop(&store);