timer: orion-timer: Add timer_get_boot_us() for BOOTSTAGE support

Add timer_get_boot_us() to support boards, that have CONFIG_BOOTSTAGE
enabled, like pogo_v4.

Signed-off-by: Stefan Roese <sr@denx.de>
Tested-by: Tony Dinh <mibodhi@gmail.com>
diff --git a/drivers/timer/orion-timer.c b/drivers/timer/orion-timer.c
index 14f318e..cd63ea9 100644
--- a/drivers/timer/orion-timer.c
+++ b/drivers/timer/orion-timer.c
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: GPL-2.0+
 #include <asm/io.h>
 #include <common.h>
+#include <div64.h>
 #include <dm/device.h>
 #include <dm/fdtaddr.h>
 #include <timer.h>
@@ -22,12 +23,56 @@
 
 #define MVEBU_TIMER_FIXED_RATE_25MHZ	25000000
 
+static bool early_init_done __section(".data") = false;
+
+/* Common functions for early (boot) and DM based timer */
+static void orion_timer_init(void *base, enum input_clock_type type)
+{
+	writel(~0, base + TIMER0_VAL);
+	writel(~0, base + TIMER0_RELOAD);
+
+	if (type == INPUT_CLOCK_25MHZ) {
+		/*
+		 * On Armada XP / 38x ..., the 25MHz clock source needs to
+		 * be enabled
+		 */
+		setbits_le32(base + TIMER_CTRL, BIT(11));
+	}
+
+	/* enable timer */
+	setbits_le32(base + TIMER_CTRL, TIMER0_EN | TIMER0_RELOAD_EN);
+}
+
+static uint64_t orion_timer_get_count(void *base)
+{
+	return timer_conv_64(~readl(base + TIMER0_VAL));
+}
+
+/* Early (e.g. bootstage etc) timer functions */
+static void notrace timer_early_init(void)
+{
+	/* Only init the timer once */
+	if (early_init_done)
+		return;
+	early_init_done = true;
+
+	if (IS_ENABLED(CONFIG_ARCH_MVEBU))
+		orion_timer_init((void *)MVEBU_TIMER_BASE, INPUT_CLOCK_25MHZ);
+	else
+		orion_timer_init((void *)MVEBU_TIMER_BASE, INPUT_CLOCK_NON_FIXED);
+}
+
 /**
  * timer_early_get_rate() - Get the timer rate before driver model
  */
 unsigned long notrace timer_early_get_rate(void)
 {
-	return MVEBU_TIMER_FIXED_RATE_25MHZ;
+	timer_early_init();
+
+	if (IS_ENABLED(CONFIG_ARCH_MVEBU))
+		return MVEBU_TIMER_FIXED_RATE_25MHZ;
+	else
+		return CONFIG_SYS_TCLK;
 }
 
 /**
@@ -36,14 +81,25 @@
  */
 u64 notrace timer_early_get_count(void)
 {
-	return timer_conv_64(~readl(MVEBU_TIMER_BASE + TIMER0_VAL));
+	timer_early_init();
+
+	return orion_timer_get_count((void *)MVEBU_TIMER_BASE);
 }
 
-static uint64_t orion_timer_get_count(struct udevice *dev)
+ulong timer_get_boot_us(void)
+{
+	u64 ticks;
+
+	ticks = timer_early_get_count();
+	return lldiv(ticks * 1000, timer_early_get_rate());
+}
+
+/* DM timer functions */
+static uint64_t dm_orion_timer_get_count(struct udevice *dev)
 {
 	struct orion_timer_priv *priv = dev_get_priv(dev);
 
-	return timer_conv_64(~readl(priv->base + TIMER0_VAL));
+	return orion_timer_get_count(priv->base);
 }
 
 static int orion_timer_probe(struct udevice *dev)
@@ -58,28 +114,17 @@
 		return -ENOMEM;
 	}
 
-	writel(~0, priv->base + TIMER0_VAL);
-	writel(~0, priv->base + TIMER0_RELOAD);
-
-	if (type == INPUT_CLOCK_25MHZ) {
-		/*
-		 * On Armada XP / 38x ..., the 25MHz clock source needs to
-		 * be enabled
-		 */
-		setbits_le32(priv->base + TIMER_CTRL, BIT(11));
+	if (type == INPUT_CLOCK_25MHZ)
 		uc_priv->clock_rate = MVEBU_TIMER_FIXED_RATE_25MHZ;
-	} else {
+	else
 		uc_priv->clock_rate = CONFIG_SYS_TCLK;
-	}
-
-	/* enable timer */
-	setbits_le32(priv->base + TIMER_CTRL, TIMER0_EN | TIMER0_RELOAD_EN);
+	orion_timer_init(priv->base, type);
 
 	return 0;
 }
 
 static const struct timer_ops orion_timer_ops = {
-	.get_count = orion_timer_get_count,
+	.get_count = dm_orion_timer_get_count,
 };
 
 static const struct udevice_id orion_timer_ids[] = {