arm/pxa: fix and cleanup of pxa_mem_setup macro v2

index 1f1759b..e2ddfe9 100644
@@ -102,11 +102,15 @@
 /*
  * This macro sets up the Memory controller of the PXA2xx CPU
  *
- * Clobbered regs: r3, r4, r5
+ * Clobbered regs: r3, r4, r5, r6, r7
+ *
+ * See section 6.4.10 of Intel PXA2xx Processor Developer's Manual
+ *   http://www.marvell.com/products/processors/applications/pxa_family/pxa_27x_dev_man.pdf
  */
-.macro	pxa_mem_setup
+.macro	pxa2xx_mem_setup
 	/* This comes handy when setting MDREFR */
 	ldr	r3, =MEMC_BASE
+	ldr	r7, =CONFIG_SYS_MDREFR_VAL
 
 	/*
 	 * 1) Initialize Asynchronous static memory controller

@@ -149,51 +153,68 @@
 	 */
 
 	/*
-	 * Before accessing MDREFR we need a valid DRI field, so we set
-	 * this to power on defaults + DRI field.
+	 * Before accessing MDREFR we need a valid DRI field.
+	 * Also we must properly configure MDREFR[K0DB2] and MDREFR[K0DB4].
+	 * Optionaly we can set MDREFR[KxFREE] bits.
+	 * So we set MDREFR to power on defaults + (DRI, K0DB2, K0DB4, KxFREE)
+	 * fields from the config.
+	 *
+	 * WARNING: K0DB2 and K0DB4 bits are usually set, while KxFREE bits
+	 *          are usually unset.
 	 */
 	ldr	r5, [r3, #MDREFR_OFFSET]
-	bic	r5, r5, #0x0ff
-	bic	r5, r5, #0xf00	/* MDREFR user config with zeroed DRI */
-
-	ldr	r4, =CONFIG_SYS_MDREFR_VAL
-	mov	r6, r4
-	lsl	r4, #20
-	lsr	r4, #20		/* Get a valid DRI field */
-
-	orr	r5, r5, r4	/* MDREFR user config with correct DRI */
+	ldr	r4, =( 0xFFF | MDREFR_K0DB4 | MDREFR_K0DB2 | \
+		       MDREFR_K0FREE | MDREFR_K1FREE | MDREFR_K2FREE )
+	bic	r5, r5, r4	/* clear DRI, K0DB2, K0DB4, KxFREE fields */
+	and	r4, r7, r4
+	orr	r5, r5, r4	/* use custom DRI, K0DB2, K0DB4, KxFREE */
 
 	orr	r5, #MDREFR_K0RUN
 	orr	r5, #MDREFR_SLFRSH
 	bic	r5, #MDREFR_APD
-	bic	r5, #MDREFR_E1PIN
+
+	/* enable them later, if SDRAM is present */
+	bic	r5, #( MDREFR_E1PIN | MDREFR_K1RUN | MDREFR_K2RUN | \
+				      MDREFR_K1DB2 | MDREFR_K2DB2 )
 
 	str	r5, [r3, #MDREFR_OFFSET]
-	ldr	r4, [r3, #MDREFR_OFFSET]
+	ldr	r5, [r3, #MDREFR_OFFSET]
 
 	/*
 	 * 5) Initialize Synchronous Static Memory (Flash/Peripherals)
 	 */
 
-	/* Initialize SXCNFG register. Assert the enable bits.
-	 *
-	 * Write SXMRS to cause an MRS command to all enabled banks of
-	 * synchronous static memory. Note that SXLCR need not be written
-	 * at this time.
+	/* Initialize SXCNFG register to enable synchronous flash memory.
+	 * While the synchronous flash banks are being configured, the SDRAM
+	 * banks must be disabled and MDREFR[APD] must be de-asserted.
 	 */
 	write32rb	(MEMC_BASE + SXCNFG_OFFSET), CONFIG_SYS_SXCNFG_VAL
 
 	/*
-	 * 6) Initialize SDRAM
+	 * 6) Initialize SDRAM,
+	 *    If SDRAM present, then MDREFR[K1RUN] and/or MDREFR[K1RUN] bits
+	 *    must be set. Also we must properly configure MDREFR[K1DB2] and
+	 *    MDREFR[K2DB2] in this case.
+	 *
+	 *    WARNING: K1DB2 and K2DB2 bits are usually set if SDRAM present
 	 */
 
+#if (CONFIG_SYS_MDREFR_VAL & (MDREFR_K1RUN | MDREFR_K2RUN))
+	and	r4, r7, #( MDREFR_K1RUN | MDREFR_K2RUN | \
+			   MDREFR_K1DB2 | MDREFR_K2DB2 )
+	ldr	r6, [r3, #MDREFR_OFFSET]
+	orr	r6, r6, r4
+	str	r6, [r3, #MDREFR_OFFSET]
+	ldr	r6, [r3, #MDREFR_OFFSET]
+
 	bic	r6, #MDREFR_SLFRSH
 	str	r6, [r3, #MDREFR_OFFSET]
-	ldr	r4, [r3, #MDREFR_OFFSET]
+	ldr	r6, [r3, #MDREFR_OFFSET]
 
 	orr	r6, #MDREFR_E1PIN
 	str	r6, [r3, #MDREFR_OFFSET]
-	ldr	r4, [r3, #MDREFR_OFFSET]
+	ldr	r6, [r3, #MDREFR_OFFSET]
+#endif
 
 	/*
 	 * 7) Write MDCNFG with MDCNFG:DEx deasserted (set to 0), to configure

@@ -209,6 +230,7 @@
 	str	r4, [r3, #MDCNFG_OFFSET]
 	ldr	r4, [r3, #MDCNFG_OFFSET]
 
+#if (CONFIG_SYS_MDREFR_VAL & (MDREFR_K1RUN | MDREFR_K2RUN))
 	/* Wait for the clock to the SDRAMs to stabilize, 100..200 usec. */
 	pxa_wait_ticks	0x300
 

@@ -226,7 +248,7 @@
 .endr
 
 	/*
-	 * 9) Write MDCNFG with enable bits asserted (MDCNFG:DEx set to 1).
+	 * 9) Set custom MDCNFG[DEx] bits to enable required SDRAM partitions
 	 */
 
 	ldr	r5, =CONFIG_SYS_MDCNFG_VAL

@@ -238,19 +260,25 @@
 	ldr     r4, [r3, #MDCNFG_OFFSET]
 
 	/*
-	 * 10) Write MDMRS.
+	 * 10) Write to MDMRS register to trigger an MRS command to
+	 *     all enabled banks of SDRAM. For each SDRAM partition pair
+	 *     that has one or both partitions enabled, this forces a pass
+	 *     through the MRS state and a return to NOP.
 	 */
 
 	ldr     r4, =CONFIG_SYS_MDMRS_VAL
 	str     r4, [r3, #MDMRS_OFFSET]
 	ldr     r4, [r3, #MDMRS_OFFSET]
+#endif
 
 	/*
-	 * 11) Enable APD
+	 * 11) Optionaly enable auto-power-down by setting MDREFR[APD]
+	 *
+	 *     WARNING: APD bit is usually set.
 	 */
 
 	ldr	r4, [r3, #MDREFR_OFFSET]
-	and	r6, r6, #MDREFR_APD
+	and	r6, r7, #MDREFR_APD
 	orr	r4, r4, r6
 	str	r4, [r3, #MDREFR_OFFSET]
 	ldr	r4, [r3, #MDREFR_OFFSET]