diff --git a/common/cmd_ubi.c b/common/cmd_ubi.c
index 10eea655701e793fd6641f2e22972b78c1651f76..0460b4cc56f352581b91c951725d99adb1597c9d 100644
--- a/common/cmd_ubi.c
+++ b/common/cmd_ubi.c
@@ -14,7 +14,7 @@
 #include <common.h>
 #include <command.h>
 #include <exports.h>
-
+#include <memalign.h>
 #include <nand.h>
 #include <onenand_uboot.h>
 #include <linux/mtd/mtd.h>
diff --git a/drivers/mtd/nand/nand_util.c b/drivers/mtd/nand/nand_util.c
index 21b4a618ce503c98bce2c861bc1e876941d83dc6..71285b6669416b1402ebd8b03f7b9957b47a37a4 100644
--- a/drivers/mtd/nand/nand_util.c
+++ b/drivers/mtd/nand/nand_util.c
@@ -23,6 +23,7 @@
 #include <command.h>
 #include <watchdog.h>
 #include <malloc.h>
+#include <memalign.h>
 #include <div64.h>
 
 #include <asm/errno.h>
diff --git a/fs/ubifs/super.c b/fs/ubifs/super.c
index 0bf52db0cef5c1ce0ecb271ab234e5cd5a174d03..41763a189790041cd311c1fd608a79c186b783e9 100644
--- a/fs/ubifs/super.c
+++ b/fs/ubifs/super.c
@@ -28,6 +28,9 @@
 #include <linux/writeback.h>
 #else
 
+#include <common.h>
+#include <malloc.h>
+#include <memalign.h>
 #include <linux/compat.h>
 #include <linux/stat.h>
 #include <linux/err.h>
diff --git a/fs/ubifs/ubifs.c b/fs/ubifs/ubifs.c
index 4daa7fad53df51acffcc7d58199181412e05b6d6..f7a084747e692ac0506cc8649b93058351e71f11 100644
--- a/fs/ubifs/ubifs.c
+++ b/fs/ubifs/ubifs.c
@@ -23,6 +23,8 @@
  *          Adrian Hunter
  */
 
+#include <common.h>
+#include <memalign.h>
 #include "ubifs.h"
 #include <u-boot/zlib.h>
 
diff --git a/include/common.h b/include/common.h
index c12f402f773ff6d0a334413437ff9bf17e03f236..c48e5bc11bf5b7ffe1d6f4314a17aa431b247b9a 100644
--- a/include/common.h
+++ b/include/common.h
@@ -1060,15 +1060,6 @@ int cpu_release(int nr, int argc, char * const argv[]);
 #define DEFINE_CACHE_ALIGN_BUFFER(type, name, size)			\
 	DEFINE_ALIGN_BUFFER(type, name, size, ARCH_DMA_MINALIGN)
 
-#ifndef __ASSEMBLY__
-#include <malloc.h>
-
-static inline void *malloc_cache_aligned(size_t size)
-{
-	return memalign(ARCH_DMA_MINALIGN, ALIGN(size, ARCH_DMA_MINALIGN));
-}
-#endif
-
 /*
  * check_member() - Check the offset of a structure member
  *
diff --git a/include/memalign.h b/include/memalign.h
new file mode 100644
index 0000000000000000000000000000000000000000..f78b9dd9dc7c31b16a0a97d97b9170bd27924f7a
--- /dev/null
+++ b/include/memalign.h
@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2015 Google, Inc
+ *
+ * SPDX-License-Identifier:    GPL-2.0+
+ */
+
+#ifndef __ALIGNMEM_H
+#define __ALIGNMEM_H
+
+/*
+ * ARCH_DMA_MINALIGN is defined in asm/cache.h for each architecture.  It
+ * is used to align DMA buffers.
+ */
+#ifndef __ASSEMBLY__
+#include <asm/cache.h>
+
+#include <malloc.h>
+
+static inline void *malloc_cache_aligned(size_t size)
+{
+	return memalign(ARCH_DMA_MINALIGN, ALIGN(size, ARCH_DMA_MINALIGN));
+}
+#endif
+
+#endif /* __ALIGNMEM_H */
diff --git a/lib/gzip.c b/lib/gzip.c
index cd8e9fea43dd6eb80f9256604cdf77b90e6b6e32..2c49e4e9ffa0c748abc972fa14263c7047259aad 100644
--- a/lib/gzip.c
+++ b/lib/gzip.c
@@ -10,6 +10,7 @@
 #include <command.h>
 #include <image.h>
 #include <malloc.h>
+#include <memalign.h>
 #include <u-boot/zlib.h>
 #include "zlib/zutil.h"
 
diff --git a/lib/zlib/zutil.c b/lib/zlib/zutil.c
index 173a81d1ea4df76325f2f45cbe82aa8e3e10641c..227343e48d3fb1305a229ac4b72559bda6351b48 100644
--- a/lib/zlib/zutil.c
+++ b/lib/zlib/zutil.c
@@ -43,7 +43,9 @@ void z_error (m)
  */
 #ifndef MY_ZCALLOC /* Any system without a special alloc function */
 
-#ifndef __UBOOT__
+#ifdef __UBOOT__
+#include <malloc.h>
+#else
 #ifndef STDC
 extern voidp    malloc OF((uInt size));
 extern voidp    calloc OF((uInt items, uInt size));