#include "head_32.h"
 
-/* 601 only have IBAT; cr0.eq is set on 601 when using this macro */
+/* 601 only have IBAT */
+#ifdef CONFIG_PPC_BOOK3S_601
+#define LOAD_BAT(n, reg, RA, RB)       \
+       li      RA,0;                   \
+       mtspr   SPRN_IBAT##n##U,RA;     \
+       lwz     RA,(n*16)+0(reg);       \
+       lwz     RB,(n*16)+4(reg);       \
+       mtspr   SPRN_IBAT##n##U,RA;     \
+       mtspr   SPRN_IBAT##n##L,RB
+#else
 #define LOAD_BAT(n, reg, RA, RB)       \
        /* see the comment for clear_bats() -- Cort */ \
        li      RA,0;                   \
        lwz     RB,(n*16)+4(reg);       \
        mtspr   SPRN_IBAT##n##U,RA;     \
        mtspr   SPRN_IBAT##n##L,RB;     \
-       beq     1f;                     \
        lwz     RA,(n*16)+8(reg);       \
        lwz     RB,(n*16)+12(reg);      \
        mtspr   SPRN_DBAT##n##U,RA;     \
-       mtspr   SPRN_DBAT##n##L,RB;     \
-1:
+       mtspr   SPRN_DBAT##n##L,RB
+#endif
 
        __HEAD
        .stabs  "arch/powerpc/kernel/",N_SO,0,0,0f
 
 /* Load the BAT registers with the values set up by MMU_init.
    MMU_init takes care of whether we're on a 601 or not. */
-       mfpvr   r3
-       srwi    r3,r3,16
-       cmpwi   r3,1
        lis     r3,BATS@ha
        addi    r3,r3,BATS@l
        tophys(r3,r3)
  */
 clear_bats:
        li      r10,0
-       mfspr   r9,SPRN_PVR
-       rlwinm  r9,r9,16,16,31          /* r9 = 1 for 601, 4 for 604 */
-       cmpwi   r9, 1
-       beq     1f
 
+#ifndef CONFIG_PPC_BOOK3S_601
        mtspr   SPRN_DBAT0U,r10
        mtspr   SPRN_DBAT0L,r10
        mtspr   SPRN_DBAT1U,r10
        mtspr   SPRN_DBAT2L,r10
        mtspr   SPRN_DBAT3U,r10
        mtspr   SPRN_DBAT3L,r10
-1:
+#endif
        mtspr   SPRN_IBAT0U,r10
        mtspr   SPRN_IBAT0L,r10
        mtspr   SPRN_IBAT1U,r10
  */
 initial_bats:
        lis     r11,PAGE_OFFSET@h
-       mfspr   r9,SPRN_PVR
-       rlwinm  r9,r9,16,16,31          /* r9 = 1 for 601, 4 for 604 */
-       cmpwi   0,r9,1
-       bne     4f
+#ifdef CONFIG_PPC_BOOK3S_601
        ori     r11,r11,4               /* set up BAT registers for 601 */
        li      r8,0x7f                 /* valid, block length = 8MB */
        mtspr   SPRN_IBAT0U,r11         /* N.B. 601 has valid bit in */
        addis   r8,r8,0x800000@h
        mtspr   SPRN_IBAT2U,r11
        mtspr   SPRN_IBAT2L,r8
-       isync
-       blr
-
-4:     tophys(r8,r11)
+#else
+       tophys(r8,r11)
 #ifdef CONFIG_SMP
        ori     r8,r8,0x12              /* R/W access, M=1 */
 #else
        mtspr   SPRN_DBAT0U,r11         /* bit in upper BAT register */
        mtspr   SPRN_IBAT0L,r8
        mtspr   SPRN_IBAT0U,r11
+#endif
        isync
        blr
 
-
 #ifdef CONFIG_BOOTX_TEXT
 setup_disp_bat:
        /*
        beqlr
        lwz     r11,0(r8)
        lwz     r8,4(r8)
-       mfspr   r9,SPRN_PVR
-       rlwinm  r9,r9,16,16,31          /* r9 = 1 for 601, 4 for 604 */
-       cmpwi   0,r9,1
-       beq     1f
+#ifndef CONFIG_PPC_BOOK3S_601
        mtspr   SPRN_DBAT3L,r8
        mtspr   SPRN_DBAT3U,r11
-       blr
-1:     mtspr   SPRN_IBAT3L,r8
+#else
+       mtspr   SPRN_IBAT3L,r8
        mtspr   SPRN_IBAT3U,r11
+#endif
        blr
 #endif /* CONFIG_BOOTX_TEXT */