[v2] cpuidle: add riscv cpuidle driver

diff --git a/arch/riscv/Kconfig b/arch/riscv/Kconfig
index df18372..799bf86 100644
--- a/arch/riscv/Kconfig
+++ b/arch/riscv/Kconfig
@@ -33,6 +33,7 @@  config RISCV
 	select ARCH_WANT_HUGE_PMD_SHARE if 64BIT
 	select CLONE_BACKWARDS
 	select COMMON_CLK
+	select CPU_IDLE
 	select EDAC_SUPPORT
 	select GENERIC_ARCH_TOPOLOGY if SMP
 	select GENERIC_ATOMIC64 if !64BIT
@@ -407,6 +408,12 @@  config BUILTIN_DTB
 	depends on RISCV_M_MODE
 	depends on OF
 
+menu "CPU Power Management"
+
+source "drivers/cpuidle/Kconfig"
+
+endmenu
+
 menu "Power management options"
 
 source "kernel/power/Kconfig"
diff --git a/arch/riscv/include/asm/cpuidle.h b/arch/riscv/include/asm/cpuidle.h
new file mode 100644
index 00000000..599b810
--- /dev/null
+++ b/arch/riscv/include/asm/cpuidle.h
@@ -0,0 +1,16 @@ 
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __RISCV_CPUIDLE_H
+#define __RISCV_CPUIDLE_H
+
+static inline void cpu_do_idle(void)
+{
+	/*
+	 * Add mb() here to ensure that all
+	 * IO/MEM access are completed prior
+	 * to enter WFI.
+	 */
+	mb();
+	wait_for_interrupt();
+}
+
+#endif
diff --git a/arch/riscv/kernel/process.c b/arch/riscv/kernel/process.c
index 2b97c49..5431aaa 100644
--- a/arch/riscv/kernel/process.c
+++ b/arch/riscv/kernel/process.c
@@ -21,6 +21,7 @@ 
 #include <asm/string.h>
 #include <asm/switch_to.h>
 #include <asm/thread_info.h>
+#include <asm/cpuidle.h>
 
 register unsigned long gp_in_global __asm__("gp");
 
@@ -35,7 +36,7 @@  extern asmlinkage void ret_from_kernel_thread(void);
 
 void arch_cpu_idle(void)
 {
-	wait_for_interrupt();
+	cpu_do_idle();
 	local_irq_enable();
 }
 
diff --git a/drivers/cpuidle/Kconfig b/drivers/cpuidle/Kconfig
index c0aeedd..f6be0fd 100644
--- a/drivers/cpuidle/Kconfig
+++ b/drivers/cpuidle/Kconfig
@@ -62,6 +62,11 @@  depends on PPC
 source "drivers/cpuidle/Kconfig.powerpc"
 endmenu
 
+menu "RISCV CPU Idle Drivers"
+depends on RISCV
+source "drivers/cpuidle/Kconfig.riscv"
+endmenu
+
 config HALTPOLL_CPUIDLE
 	tristate "Halt poll cpuidle driver"
 	depends on X86 && KVM_GUEST
diff --git a/drivers/cpuidle/Kconfig.riscv b/drivers/cpuidle/Kconfig.riscv
new file mode 100644
index 00000000..7bec059
--- /dev/null
+++ b/drivers/cpuidle/Kconfig.riscv
@@ -0,0 +1,11 @@ 
+# SPDX-License-Identifier: GPL-2.0-only
+#
+# RISCV CPU Idle drivers
+#
+config RISCV_CPUIDLE
+	bool "Generic RISCV CPU idle Driver"
+	select DT_IDLE_STATES
+	select CPU_IDLE_MULTIPLE_DRIVERS
+	help
+	  Select this option to enable generic cpuidle driver for RISCV.
+	  Now only support C0 State.
diff --git a/drivers/cpuidle/Makefile b/drivers/cpuidle/Makefile
index 26bbc5e..4c83c4e 100644
--- a/drivers/cpuidle/Makefile
+++ b/drivers/cpuidle/Makefile
@@ -34,3 +34,7 @@  obj-$(CONFIG_MIPS_CPS_CPUIDLE)		+= cpuidle-cps.o
 # POWERPC drivers
 obj-$(CONFIG_PSERIES_CPUIDLE)		+= cpuidle-pseries.o
 obj-$(CONFIG_POWERNV_CPUIDLE)		+= cpuidle-powernv.o
+
+###############################################################################
+# RISCV drivers
+obj-$(CONFIG_RISCV_CPUIDLE)		+= cpuidle-riscv.o
diff --git a/drivers/cpuidle/cpuidle-riscv.c b/drivers/cpuidle/cpuidle-riscv.c
new file mode 100644
index 00000000..2ada1d5
--- /dev/null
+++ b/drivers/cpuidle/cpuidle-riscv.c
@@ -0,0 +1,54 @@ 
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * RISC-V CPU idle driver.
+ *
+ * Copyright (C) 2020-2022 Allwinner Ltd
+ *
+ * Based on code - driver/cpuidle/cpuidle-at91.c
+ *
+ */
+#include <linux/cpuidle.h>
+#include <linux/cpumask.h>
+#include <linux/cpu_pm.h>
+#include <linux/kernel.h>
+#include <linux/module.h>
+#include <linux/of.h>
+#include <linux/slab.h>
+#include <linux/platform_device.h>
+#include <asm/cpuidle.h>
+
+#define MAX_IDLE_STATES	1
+
+int __weak riscv_low_level_suspend_enter(int state)
+{
+	return 0;
+}
+
+/* Actual code that puts the SoC in different idle states */
+static int riscv_enter_idle(struct cpuidle_device *dev,
+			struct cpuidle_driver *drv,
+			       int index)
+{
+	return CPU_PM_CPU_IDLE_ENTER_PARAM(riscv_low_level_suspend_enter,
+					   index, 0);
+}
+
+static struct cpuidle_driver riscv_idle_driver = {
+	.name			= "riscv_idle",
+	.owner			= THIS_MODULE,
+	.states[0]		= {
+		.enter			= riscv_enter_idle,
+		.exit_latency		= 1,
+		.target_residency	= 1,
+		.name			= "WFI",
+		.desc			= "RISCV WFI",
+	},
+	.state_count = MAX_IDLE_STATES,
+};
+
+static int __init riscv_cpuidle_init(void)
+{
+	return cpuidle_register(&riscv_idle_driver, NULL);
+}
+
+device_initcall(riscv_cpuidle_init);