diff --git a/arch/arm/cpu/armv7/omap-common/mem-common.c b/arch/arm/cpu/armv7/omap-common/mem-common.c
index 136a032..d72e82e 100644
--- a/arch/arm/cpu/armv7/omap-common/mem-common.c
+++ b/arch/arm/cpu/armv7/omap-common/mem-common.c
@@ -20,9 +20,20 @@
 #include <asm/arch/sys_proto.h>
 #include <command.h>
 #include <linux/mtd/omap_gpmc.h>
+#include <jffs2/load_kernel.h>
 
 const struct gpmc *gpmc_cfg = (struct gpmc *)GPMC_BASE;
 
+#if defined(CONFIG_NOR)
+char gpmc_cs0_flash = MTD_DEV_TYPE_NOR;
+#elif defined(CONFIG_NAND) || defined(CONFIG_CMD_NAND)
+char gpmc_cs0_flash = MTD_DEV_TYPE_NAND;
+#elif defined(CONFIG_CMD_ONENAND)
+char gpmc_cs0_flash = MTD_DEV_TYPE_ONENAND;
+#else
+char gpmc_cs0_flash = -1;
+#endif
+
 #if defined(CONFIG_OMAP34XX)
 /********************************************************
  *  mem_ok() - test used to see if timings are correct
@@ -68,6 +79,81 @@
 	sdelay(2000);
 }
 
+void set_gpmc_cs0(int flash_type)
+{
+	const u32 *gpmc_regs;
+	u32 base, size;
+#if defined(CONFIG_NOR)
+	const u32 gpmc_regs_nor[GPMC_MAX_REG] = {
+		STNOR_GPMC_CONFIG1,
+		STNOR_GPMC_CONFIG2,
+		STNOR_GPMC_CONFIG3,
+		STNOR_GPMC_CONFIG4,
+		STNOR_GPMC_CONFIG5,
+		STNOR_GPMC_CONFIG6,
+		STNOR_GPMC_CONFIG7
+	};
+#endif
+#if defined(CONFIG_NAND) || defined(CONFIG_CMD_NAND)
+	const u32 gpmc_regs_nand[GPMC_MAX_REG] = {
+		M_NAND_GPMC_CONFIG1,
+		M_NAND_GPMC_CONFIG2,
+		M_NAND_GPMC_CONFIG3,
+		M_NAND_GPMC_CONFIG4,
+		M_NAND_GPMC_CONFIG5,
+		M_NAND_GPMC_CONFIG6,
+		0
+	};
+#endif
+#if defined(CONFIG_CMD_ONENAND)
+	const u32 gpmc_regs_onenand[GPMC_MAX_REG] = {
+		ONENAND_GPMC_CONFIG1,
+		ONENAND_GPMC_CONFIG2,
+		ONENAND_GPMC_CONFIG3,
+		ONENAND_GPMC_CONFIG4,
+		ONENAND_GPMC_CONFIG5,
+		ONENAND_GPMC_CONFIG6,
+		0
+	};
+#endif
+
+	switch (flash_type) {
+#if defined(CONFIG_NOR)
+	case MTD_DEV_TYPE_NOR:
+		gpmc_regs = gpmc_regs_nor;
+		base = CONFIG_SYS_FLASH_BASE;
+		size = (CONFIG_SYS_FLASH_SIZE > 0x08000000) ? GPMC_SIZE_256M :
+		      ((CONFIG_SYS_FLASH_SIZE > 0x04000000) ? GPMC_SIZE_128M :
+		      ((CONFIG_SYS_FLASH_SIZE > 0x02000000) ? GPMC_SIZE_64M  :
+		      ((CONFIG_SYS_FLASH_SIZE > 0x01000000) ? GPMC_SIZE_32M  :
+		                                              GPMC_SIZE_16M)));
+		break;
+#endif
+#if defined(CONFIG_NAND) || defined(CONFIG_CMD_NAND)
+	case MTD_DEV_TYPE_NAND:
+		gpmc_regs = gpmc_regs_nand;
+		base = CONFIG_SYS_NAND_BASE;
+		size = GPMC_SIZE_16M;
+		break;
+#endif
+#if defined(CONFIG_CMD_ONENAND)
+	case MTD_DEV_TYPE_ONENAND:
+		gpmc_regs = gpmc_regs_onenand;
+		base = CONFIG_SYS_ONENAND_BASE;
+		size = GPMC_SIZE_128M;
+		break;
+#endif
+	default:
+		/* disable the GPMC0 config set by ROM code */
+		writel(0, &gpmc_cfg->cs[0].config7);
+		sdelay(1000);
+		return;
+	}
+
+	/* enable chip-select specific configurations */
+	enable_gpmc_cs_config(gpmc_regs, &gpmc_cfg->cs[0], base, size);
+}
+
 /*****************************************************
  * gpmc_init(): init gpmc bus
  * Init GPMC for x16, MuxMode (SDRAM in x32).
@@ -75,68 +161,14 @@
  *****************************************************/
 void gpmc_init(void)
 {
-#if defined(CONFIG_NOR)
-/* configure GPMC for NOR */
-	const u32 gpmc_regs[GPMC_MAX_REG] = {	STNOR_GPMC_CONFIG1,
-						STNOR_GPMC_CONFIG2,
-						STNOR_GPMC_CONFIG3,
-						STNOR_GPMC_CONFIG4,
-						STNOR_GPMC_CONFIG5,
-						STNOR_GPMC_CONFIG6,
-						STNOR_GPMC_CONFIG7
-						};
-	u32 base = CONFIG_SYS_FLASH_BASE;
-	u32 size =	(CONFIG_SYS_FLASH_SIZE  > 0x08000000) ? GPMC_SIZE_256M :
-	/* > 64MB */	((CONFIG_SYS_FLASH_SIZE > 0x04000000) ? GPMC_SIZE_128M :
-	/* > 32MB */	((CONFIG_SYS_FLASH_SIZE > 0x02000000) ? GPMC_SIZE_64M  :
-	/* > 16MB */	((CONFIG_SYS_FLASH_SIZE > 0x01000000) ? GPMC_SIZE_32M  :
-	/* min 16MB */	GPMC_SIZE_16M)));
-#elif defined(CONFIG_NAND) || defined(CONFIG_CMD_NAND)
-/* configure GPMC for NAND */
-	const u32  gpmc_regs[GPMC_MAX_REG] = {	M_NAND_GPMC_CONFIG1,
-						M_NAND_GPMC_CONFIG2,
-						M_NAND_GPMC_CONFIG3,
-						M_NAND_GPMC_CONFIG4,
-						M_NAND_GPMC_CONFIG5,
-						M_NAND_GPMC_CONFIG6,
-						0
-						};
-	u32 base = CONFIG_SYS_NAND_BASE;
-	u32 size = GPMC_SIZE_16M;
-
-#elif defined(CONFIG_CMD_ONENAND)
-	const u32 gpmc_regs[GPMC_MAX_REG] = {	ONENAND_GPMC_CONFIG1,
-						ONENAND_GPMC_CONFIG2,
-						ONENAND_GPMC_CONFIG3,
-						ONENAND_GPMC_CONFIG4,
-						ONENAND_GPMC_CONFIG5,
-						ONENAND_GPMC_CONFIG6,
-						0
-						};
-	u32 size = GPMC_SIZE_128M;
-	u32 base = CONFIG_SYS_ONENAND_BASE;
-#else
-	const u32 gpmc_regs[GPMC_MAX_REG] = { 0, 0, 0, 0, 0, 0, 0 };
-	u32 size = 0;
-	u32 base = 0;
-#endif
 	/* global settings */
 	writel(0x00000008, &gpmc_cfg->sysconfig);
 	writel(0x00000000, &gpmc_cfg->irqstatus);
 	writel(0x00000000, &gpmc_cfg->irqenable);
 	/* disable timeout, set a safe reset value */
 	writel(0x00001ff0, &gpmc_cfg->timeout_control);
-#ifdef CONFIG_NOR
-	writel(0x00000200, &gpmc_cfg->config);
-#else
-	writel(0x00000012, &gpmc_cfg->config);
-#endif
-	/*
-	 * Disable the GPMC0 config set by ROM code
-	 */
-	writel(0, &gpmc_cfg->cs[0].config7);
-	sdelay(1000);
-	/* enable chip-select specific configurations */
-	if (base != 0)
-		enable_gpmc_cs_config(gpmc_regs, &gpmc_cfg->cs[0], base, size);
+	writel(gpmc_cs0_flash == MTD_DEV_TYPE_NOR ?
+		0x00000200 : 0x00000012, &gpmc_cfg->config);
+
+	set_gpmc_cs0(gpmc_cs0_flash);
 }