diff --git a/common/cmd_ubi.c b/common/cmd_ubi.c
index 10eea65..0460b4c 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 21b4a61..71285b6 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 0bf52db..41763a1 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 4daa7fa..f7a0847 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 c12f402..c48e5bc 100644
--- a/include/common.h
+++ b/include/common.h
@@ -1060,15 +1060,6 @@
 #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 0000000..f78b9dd
--- /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 cd8e9fe..2c49e4e 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 173a81d..227343e 100644
--- a/lib/zlib/zutil.c
+++ b/lib/zlib/zutil.c
@@ -43,7 +43,9 @@
  */
 #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));