/*
  * Copyright (c) 2011 Zhang, Keguang <keguang.zhang@gmail.com>
  *
- * Loongson 1 watchdog register definitions.
+ * Loongson 1 Watchdog Register Definitions.
  *
  * This program is free software; you can redistribute it and/or modify it
  * under  the terms of the GNU General  Public License as published by the
 #ifndef __ASM_MACH_LOONGSON1_REGS_WDT_H
 #define __ASM_MACH_LOONGSON1_REGS_WDT_H
 
-#define LS1X_WDT_REG(x) \
-               ((void __iomem *)KSEG1ADDR(LS1X_WDT_BASE + (x)))
-
-#define LS1X_WDT_EN                    LS1X_WDT_REG(0x0)
-#define LS1X_WDT_SET                   LS1X_WDT_REG(0x4)
-#define LS1X_WDT_TIMER                 LS1X_WDT_REG(0x8)
+#define WDT_EN                 0x0
+#define WDT_TIMER              0x4
+#define WDT_SET                        0x8
 
 #endif /* __ASM_MACH_LOONGSON1_REGS_WDT_H */
 
 
 #include <loongson1.h>
 
-static void ls1x_restart(char *command)
-{
-       __raw_writel(0x1, LS1X_WDT_EN);
-       __raw_writel(0x5000000, LS1X_WDT_TIMER);
-       __raw_writel(0x1, LS1X_WDT_SET);
-}
+static void __iomem *wdt_base;
 
 static void ls1x_halt(void)
 {
        }
 }
 
+static void ls1x_restart(char *command)
+{
+       __raw_writel(0x1, wdt_base + WDT_EN);
+       __raw_writel(0x1, wdt_base + WDT_TIMER);
+       __raw_writel(0x1, wdt_base + WDT_SET);
+
+       ls1x_halt();
+}
+
 static void ls1x_power_off(void)
 {
        ls1x_halt();
 
 static int __init ls1x_reboot_setup(void)
 {
+       wdt_base = ioremap_nocache(LS1X_WDT_BASE, 0x0f);
+       if (!wdt_base)
+               panic("Failed to remap watchdog registers");
+
        _machine_restart = ls1x_restart;
        _machine_halt = ls1x_halt;
        pm_power_off = ls1x_power_off;