#include <linux/err.h>
 #include <linux/of.h>
 
-#include <asm/mach-jz4740/timer.h>
-
-#define JZ_WDT_CLOCK_PCLK 0x1
-#define JZ_WDT_CLOCK_RTC  0x2
-#define JZ_WDT_CLOCK_EXT  0x4
-
-#define JZ_WDT_CLOCK_DIV_1    (0 << TCU_TCSR_PRESCALE_LSB)
-#define JZ_WDT_CLOCK_DIV_4    (1 << TCU_TCSR_PRESCALE_LSB)
-#define JZ_WDT_CLOCK_DIV_16   (2 << TCU_TCSR_PRESCALE_LSB)
-#define JZ_WDT_CLOCK_DIV_64   (3 << TCU_TCSR_PRESCALE_LSB)
-#define JZ_WDT_CLOCK_DIV_256  (4 << TCU_TCSR_PRESCALE_LSB)
-#define JZ_WDT_CLOCK_DIV_1024 (5 << TCU_TCSR_PRESCALE_LSB)
-
 #define DEFAULT_HEARTBEAT 5
 #define MAX_HEARTBEAT     2048
 
 struct jz4740_wdt_drvdata {
        struct watchdog_device wdt;
        void __iomem *base;
-       struct clk *rtc_clk;
+       struct clk *clk;
+       unsigned long clk_rate;
 };
 
 static int jz4740_wdt_ping(struct watchdog_device *wdt_dev)
                                    unsigned int new_timeout)
 {
        struct jz4740_wdt_drvdata *drvdata = watchdog_get_drvdata(wdt_dev);
-       unsigned int rtc_clk_rate;
-       unsigned int timeout_value;
-       unsigned short clock_div = JZ_WDT_CLOCK_DIV_1;
+       u16 timeout_value = (u16)(drvdata->clk_rate * new_timeout);
        u8 tcer;
 
-       rtc_clk_rate = clk_get_rate(drvdata->rtc_clk);
-
-       timeout_value = rtc_clk_rate * new_timeout;
-       while (timeout_value > 0xffff) {
-               if (clock_div == JZ_WDT_CLOCK_DIV_1024) {
-                       /* Requested timeout too high;
-                       * use highest possible value. */
-                       timeout_value = 0xffff;
-                       break;
-               }
-               timeout_value >>= 2;
-               clock_div += (1 << TCU_TCSR_PRESCALE_LSB);
-       }
-
        tcer = readb(drvdata->base + TCU_REG_WDT_TCER);
        writeb(0x0, drvdata->base + TCU_REG_WDT_TCER);
-       writew(clock_div, drvdata->base + TCU_REG_WDT_TCSR);
 
        writew((u16)timeout_value, drvdata->base + TCU_REG_WDT_TDR);
        writew(0x0, drvdata->base + TCU_REG_WDT_TCNT);
-       writew(clock_div | JZ_WDT_CLOCK_RTC, drvdata->base + TCU_REG_WDT_TCSR);
 
        if (tcer & TCU_WDT_TCER_TCEN)
                writeb(TCU_WDT_TCER_TCEN, drvdata->base + TCU_REG_WDT_TCER);
 static int jz4740_wdt_start(struct watchdog_device *wdt_dev)
 {
        struct jz4740_wdt_drvdata *drvdata = watchdog_get_drvdata(wdt_dev);
+       int ret;
        u8 tcer;
 
+       ret = clk_prepare_enable(drvdata->clk);
+       if (ret)
+               return ret;
+
        tcer = readb(drvdata->base + TCU_REG_WDT_TCER);
 
-       jz4740_timer_enable_watchdog();
        jz4740_wdt_set_timeout(wdt_dev, wdt_dev->timeout);
 
        /* Start watchdog if it wasn't started already */
        struct jz4740_wdt_drvdata *drvdata = watchdog_get_drvdata(wdt_dev);
 
        writeb(0x0, drvdata->base + TCU_REG_WDT_TCER);
-       jz4740_timer_disable_watchdog();
+       clk_disable_unprepare(drvdata->clk);
 
        return 0;
 }
        struct device *dev = &pdev->dev;
        struct jz4740_wdt_drvdata *drvdata;
        struct watchdog_device *jz4740_wdt;
+       long rate;
+       int ret;
 
        drvdata = devm_kzalloc(dev, sizeof(struct jz4740_wdt_drvdata),
                               GFP_KERNEL);
        if (!drvdata)
                return -ENOMEM;
 
-       if (heartbeat < 1 || heartbeat > MAX_HEARTBEAT)
-               heartbeat = DEFAULT_HEARTBEAT;
+       drvdata->clk = devm_clk_get(&pdev->dev, "wdt");
+       if (IS_ERR(drvdata->clk)) {
+               dev_err(&pdev->dev, "cannot find WDT clock\n");
+               return PTR_ERR(drvdata->clk);
+       }
+
+       /* Set smallest clock possible */
+       rate = clk_round_rate(drvdata->clk, 1);
+       if (rate < 0)
+               return rate;
+
+       ret = clk_set_rate(drvdata->clk, rate);
+       if (ret)
+               return ret;
 
+       drvdata->clk_rate = rate;
        jz4740_wdt = &drvdata->wdt;
        jz4740_wdt->info = &jz4740_wdt_info;
        jz4740_wdt->ops = &jz4740_wdt_ops;
-       jz4740_wdt->timeout = heartbeat;
        jz4740_wdt->min_timeout = 1;
-       jz4740_wdt->max_timeout = MAX_HEARTBEAT;
+       jz4740_wdt->max_timeout = 0xffff / rate;
+       jz4740_wdt->timeout = clamp(heartbeat,
+                                   jz4740_wdt->min_timeout,
+                                   jz4740_wdt->max_timeout);
        jz4740_wdt->parent = dev;
        watchdog_set_nowayout(jz4740_wdt, nowayout);
        watchdog_set_drvdata(jz4740_wdt, drvdata);
        if (IS_ERR(drvdata->base))
                return PTR_ERR(drvdata->base);
 
-       drvdata->rtc_clk = devm_clk_get(dev, "rtc");
-       if (IS_ERR(drvdata->rtc_clk)) {
-               dev_err(dev, "cannot find RTC clock\n");
-               return PTR_ERR(drvdata->rtc_clk);
-       }
-
        return devm_watchdog_register_device(dev, &drvdata->wdt);
 }