edgetpu: sync with maran kernel

Squash of the following commits:

d81075b34256 misc: edgetpu: Add tracing for memory mapping operations
d8855fbee4b5 misc: edgetpu: Add debugfs interface to adjust TPU rails
8c187ca97eab misc: edgetpu: Expand available debugfs functions
a36bcf16d0f1 misc: edgetpu: Add policy_set API to throttle tpu
7b191e6b0537 misc:edgetpu: Initlal commit for thermal driver
f7056a62542b misc:edgetpu: Refactor Thermal driver
01720f335e86 misc: edgetpu: refactor ACPM DVFS definitions

Plus #ifdefs fixes for unit tests

Signed-off-by: Erick Reyes <erickreyes@google.com>
Change-Id: I680704b23d1786b735eef58516c009d094fc3a33

8 files changed, 630 insertions, 54 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 755b7be..8610669 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -3,10 +3,11 @@ ccflags-y	+= -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -I$(src)/include
 
 edgetpu-fw-objs := edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-shared-fw.o
 edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-map-dmabuf.o edgetpu-async.o $(edgetpu-fw-objs)
-abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-direct.o abrolhos-core.o abrolhos-platform.o abrolhos-iommu.o abrolhos-firmware.o $(edgetpu-objs)
+abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-direct.o abrolhos-core.o abrolhos-platform.o abrolhos-iommu.o abrolhos-firmware.o edgetpu-thermal.o $(edgetpu-objs)
 CFLAGS_abrolhos-direct.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-core.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device-group.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-firmware.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-platform.o := -DCONFIG_ABROLHOS=1
+CFLAGS_edgetpu-thermal.o := -DCONFIG_ABROLHOS=1
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index a88c15d..111987b 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -8,7 +8,7 @@ obj-$(CONFIG_ABROLHOS)		+= abrolhos.o
 edgetpu-fw-objs := edgetpu-firmware-util.o edgetpu-shared-fw.o edgetpu-firmware.o
 edgetpu-objs	:= edgetpu-core.o edgetpu-mailbox.o edgetpu-kci.o edgetpu-device-group.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-map-dmabuf.o edgetpu-async.o $(edgetpu-fw-objs)
 
-abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-direct.o abrolhos-platform.o abrolhos-iommu.o $(edgetpu-objs)
+abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-direct.o abrolhos-platform.o abrolhos-iommu.o edgetpu-thermal.o $(edgetpu-objs)
 
 KBUILD_OPTIONS += CONFIG_ABROLHOS=m
 
diff --git a/drivers/edgetpu/abrolhos-platform.c b/drivers/edgetpu/abrolhos-platform.c
index 921fffb..c2690ed 100644
--- a/drivers/edgetpu/abrolhos-platform.c
+++ b/drivers/edgetpu/abrolhos-platform.c
@@ -32,19 +32,10 @@
 static int power_state = TPU_ACTIVE_NOM;
 module_param(power_state, int, 0660);
 
-#define CLK_CON_DIV_CLKCMU_TPU_TPUCTL		0x1e0818f0
-#define CLK_CON_DIV_CLKCMU_TPU_BUS		0x1e0818e8
+#define TPU_INIT_FREQ		0
 
 static struct dentry *edgetpu_pwr_debugfs_dir;
 
-#if IS_ENABLED(CONFIG_ACPM_DVFS)
-
-#define TPU_DOMAIN		14
-#define TPU_INIT_FREQ		0
-extern int exynos_acpm_set_rate(unsigned int id, unsigned long rate);
-extern int exynos_acpm_set_init_freq(unsigned int dfs_id, unsigned long freq);
-extern unsigned long exynos_acpm_get_rate(unsigned int id);
-
 static int edgetpu_pwr_state_init(struct device *dev)
 {
 	int ret;
@@ -56,7 +47,7 @@ static int edgetpu_pwr_state_init(struct device *dev)
 		return ret;
 	}
 
-	ret = exynos_acpm_set_init_freq(TPU_DOMAIN, TPU_INIT_FREQ);
+	ret = exynos_acpm_set_init_freq(TPU_ACPM_DOMAIN, TPU_INIT_FREQ);
 	if (ret) {
 		dev_err(dev, "error initializing tpu state: %d\n", ret);
 		return ret;
@@ -77,7 +68,7 @@ static int edgetpu_pwr_state_set(void *data, u64 val)
 	int curr_state;
 	struct device *dev = (struct device *)data;
 
-	curr_state = exynos_acpm_get_rate(TPU_DOMAIN);
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 
 	if (curr_state == TPU_OFF && val > TPU_OFF) {
 		ret = pm_runtime_get_sync(dev);
@@ -87,7 +78,7 @@ static int edgetpu_pwr_state_set(void *data, u64 val)
 		}
 	}
 
-	ret = exynos_acpm_set_rate(TPU_DOMAIN, (unsigned long)val);
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, (unsigned long)val);
 	if (ret) {
 		dev_err(dev, "error setting tpu state: %d\n", ret);
 		return ret;
@@ -108,89 +99,226 @@ static int edgetpu_pwr_state_get(void *data, u64 *val)
 {
 	struct device *dev = (struct device *)data;
 
-	*val = exynos_acpm_get_rate(TPU_DOMAIN);
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 	dev_info(dev, "current tpu state: %d\n", *val);
 
 	return 0;
 }
 
-#else /* !CONFIG_ACPM_DVFS */
+int edgetpu_pwr_policy_set(void *data, u64 val)
+{
+	struct edgetpu_platform_dev *edgetpu_pdev =
+		(struct edgetpu_platform_dev *)data;
+	struct edgetpu_platform_pwr *platform_pwr =
+		&edgetpu_pdev->platform_pwr;
+	int ret;
 
-static int edgetpu_pwr_state_init(struct device *dev)
+	mutex_lock(&platform_pwr->policy_lock);
+	ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, val);
+
+	if (ret) {
+		dev_err(edgetpu_pdev->edgetpu_dev.dev,
+				"unable to set policy %lld (ret %d)\n",
+				val, ret);
+		mutex_unlock(&platform_pwr->policy_lock);
+		return ret;
+	}
+
+	platform_pwr->curr_policy = val;
+	mutex_unlock(&platform_pwr->policy_lock);
+	return 0;
+}
+
+static int edgetpu_pwr_policy_get(void *data, u64 *val)
 {
+	struct edgetpu_platform_dev *edgetpu_pdev =
+		(struct edgetpu_platform_dev *)data;
+	struct edgetpu_platform_pwr *platform_pwr =
+		&edgetpu_pdev->platform_pwr;
+
+	mutex_lock(&platform_pwr->policy_lock);
+	*val = platform_pwr->curr_policy;
+	mutex_unlock(&platform_pwr->policy_lock);
+
 	return 0;
 }
 
-static int edgetpu_pwr_state_set(void *data, u64 val)
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_policy, edgetpu_pwr_policy_get,
+		edgetpu_pwr_policy_set, "%llu\n");
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, edgetpu_pwr_state_get,
+		edgetpu_pwr_state_set, "%llu\n");
+
+static int edgetpu_core_rate_get(void *data, u64 *val)
 {
-	return -ENODEV;
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG);
+	return 0;
 }
 
-static int edgetpu_pwr_state_get(void *data, u64 *val)
+static int edgetpu_core_rate_set(void *data, u64 val)
 {
-	return -ENODEV;
+	int ret;
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG;
+	dbg_rate_req |= (val / 1000);
+
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
 }
 
-#endif /* CONFIG_ACPM_DVFS */
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_core_rate, edgetpu_core_rate_get,
+		edgetpu_core_rate_set, "%llu\n");
 
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, edgetpu_pwr_state_get,
-		edgetpu_pwr_state_set, "%llu\n");
+static int edgetpu_ctl_rate_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG);
+	return 0;
+}
 
-static int edgetpu_divclk_tpuctl_get(void *data, u64 *val)
+static int edgetpu_ctl_rate_set(void *data, u64 val)
 {
-	struct device *dev = (struct device *)data;
-	void __iomem *divclk_tpuctl = ioremap(CLK_CON_DIV_CLKCMU_TPU_TPUCTL, 4);
+	int ret;
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG;
+	dbg_rate_req |= (val / 1000);
 
-	*val = readl(divclk_tpuctl);
-	dev_info(dev, "divclk_tpuctl val: 0x%llx ,div ratio:%lld",
-			*val, (*val) & 0xf);
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
+}
 
-	iounmap(divclk_tpuctl);
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_ctl_rate, edgetpu_ctl_rate_get,
+		edgetpu_ctl_rate_set, "%llu\n");
 
+static int edgetpu_axi_rate_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG);
 	return 0;
 }
 
-static int edgetpu_divclk_tpuctl_set(void *data, u64 val)
+static int edgetpu_axi_rate_set(void *data, u64 val)
 {
-	void __iomem *divclk_tpuctl = ioremap(CLK_CON_DIV_CLKCMU_TPU_TPUCTL, 4);
+	int ret;
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG;
+	dbg_rate_req |= (val / 1000);
 
-	writel(val, divclk_tpuctl);
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
+}
 
-	iounmap(divclk_tpuctl);
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_axi_rate, edgetpu_axi_rate_get,
+		edgetpu_axi_rate_set, "%llu\n");
 
+
+static int edgetpu_apb_rate_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_CLK_APB_DEBUG);
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_apb_rate, edgetpu_apb_rate_get,
+		NULL, "%llu\n");
+
+static int edgetpu_uart_rate_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_CLK_UART_DEBUG);
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_uart_rate, edgetpu_uart_rate_get,
+		NULL, "%llu\n");
+
+static int edgetpu_vdd_int_m_set(void *data, u64 val)
+{
+	int ret;
+	struct device *dev = (struct device *)data;
+	unsigned long dbg_rate_req;
+
+	if (val > 1250000) {
+		dev_err(dev, "Preventing INT_M voltage > 1250000uV");
+		return -EINVAL;
+	}
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG;
+	dbg_rate_req |= val;
+
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
+}
+
+static int edgetpu_vdd_int_m_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG);
 	return 0;
 }
 
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_divclk_tpuctl, edgetpu_divclk_tpuctl_get,
-		edgetpu_divclk_tpuctl_set, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_int_m, edgetpu_vdd_int_m_get,
+		edgetpu_vdd_int_m_set, "%llu\n");
 
-static int edgetpu_divclk_tpubus_get(void *data, u64 *val)
+static int edgetpu_vdd_tpu_set(void *data, u64 val)
 {
+	int ret;
 	struct device *dev = (struct device *)data;
-	void __iomem *divclk_tpubus = ioremap(CLK_CON_DIV_CLKCMU_TPU_BUS, 4);
+	unsigned long dbg_rate_req;
 
-	*val = readl(divclk_tpubus);
-	dev_info(dev, "divclk_tpubus val: 0x%llx ,div ratio:%lld",
-			*val, (*val) & 0xf);
+	if (val > 1250000) {
+		dev_err(dev, "Preventing VDD_TPU voltage > 1250000uV");
+		return -EINVAL;
+	}
 
-	iounmap(divclk_tpubus);
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG;
+	dbg_rate_req |= val;
 
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
+}
+
+static int edgetpu_vdd_tpu_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG);
 	return 0;
 }
 
-static int edgetpu_divclk_tpubus_set(void *data, u64 val)
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu, edgetpu_vdd_tpu_get,
+		edgetpu_vdd_tpu_set, "%llu\n");
+
+static int edgetpu_vdd_tpu_m_set(void *data, u64 val)
 {
-	void __iomem *divclk_tpubus = ioremap(CLK_CON_DIV_CLKCMU_TPU_BUS, 4);
+	int ret;
+	struct device *dev = (struct device *)data;
+	unsigned long dbg_rate_req;
 
-	writel(val, divclk_tpubus);
+	if (val > 1250000) {
+		dev_err(dev, "Preventing VDD_TPU voltage > 1250000uV");
+		return -EINVAL;
+	}
 
-	iounmap(divclk_tpubus);
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG;
+	dbg_rate_req |= val;
+
+	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, dbg_rate_req);
+	return ret;
+}
 
+static int edgetpu_vdd_tpu_m_get(void *data, u64 *val)
+{
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
+			TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG);
 	return 0;
 }
 
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_divclk_tpubus, edgetpu_divclk_tpubus_get,
-		edgetpu_divclk_tpubus_set, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu_m, edgetpu_vdd_tpu_m_get,
+		edgetpu_vdd_tpu_m_set, "%llu\n");
 
 static const struct of_device_id edgetpu_of_match[] = {
 	{ .compatible = "google,darwinn", },
@@ -378,6 +506,9 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 	edgetpu_pdev->edgetpu_dev.mcp_id = -1;
 	edgetpu_pdev->edgetpu_dev.mcp_die_index = 0;
 
+	mutex_init(&edgetpu_pdev->platform_pwr.policy_lock);
+	edgetpu_pdev->platform_pwr.curr_policy = TPU_POLICY_MAX;
+
 	ret = edgetpu_device_add(&edgetpu_pdev->edgetpu_dev, &regs);
 	if (!ret && edgetpu_pdev->irq >= 0)
 		ret = edgetpu_register_irq(&edgetpu_pdev->edgetpu_dev,
@@ -425,14 +556,31 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 			  "run default firmware %s failed: %d\n",
 			  EDGETPU_DEFAULT_FIRMWARE_NAME, ret);
 
+
+	edgetpu_pdev->edgetpu_dev.thermal = devm_tpu_thermal_create(dev);
+
 	edgetpu_pwr_debugfs_dir =
 		debugfs_create_dir("edgetpu_pwr_state", NULL);
 	debugfs_create_file("pwr_state", 0660, edgetpu_pwr_debugfs_dir,
 			dev, &fops_tpu_pwr_state);
-	debugfs_create_file("tpu_divclk_tpuctl", 0660, edgetpu_pwr_debugfs_dir,
-			dev, &fops_tpu_divclk_tpuctl);
-	debugfs_create_file("tpu_divclk_tpubus", 0660, edgetpu_pwr_debugfs_dir,
-			dev, &fops_tpu_divclk_tpubus);
+	debugfs_create_file("vdd_tpu", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_vdd_tpu);
+	debugfs_create_file("vdd_tpu_m", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_vdd_tpu_m);
+	debugfs_create_file("vdd_int_m", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_vdd_int_m);
+	debugfs_create_file("core_rate", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_core_rate);
+	debugfs_create_file("ctl_rate", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_ctl_rate);
+	debugfs_create_file("axi_rate", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_axi_rate);
+	debugfs_create_file("apb_rate", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_apb_rate);
+	debugfs_create_file("uart_rate", 0660, edgetpu_pwr_debugfs_dir,
+			dev, &fops_tpu_uart_rate);
+	debugfs_create_file("policy", 0660, edgetpu_pwr_debugfs_dir,
+			edgetpu_pdev, &fops_tpu_pwr_policy);
 
 	return 0;
 }
diff --git a/drivers/edgetpu/abrolhos-platform.h b/drivers/edgetpu/abrolhos-platform.h
index b9f9430..d4d7f82 100644
--- a/drivers/edgetpu/abrolhos-platform.h
+++ b/drivers/edgetpu/abrolhos-platform.h
@@ -8,9 +8,92 @@
 #define __EDGETPU_PLATFORM_H__
 
 #include "edgetpu-internal.h"
+#include "abrolhos-pwr.h"
+
+#define TPU_ACPM_DOMAIN			14
+
+#define TPU_DEBUG_REQ			(1 << 31)
+#define TPU_VDD_TPU_DEBUG		(0 << 28)
+#define TPU_VDD_TPU_M_DEBUG		(1 << 28)
+#define TPU_VDD_INT_M_DEBUG		(2 << 28)
+#define TPU_CLK_CORE_DEBUG		(3 << 28)
+#define TPU_CLK_CTL_DEBUG		(4 << 28)
+#define TPU_CLK_AXI_DEBUG		(5 << 28)
+#define TPU_CLK_APB_DEBUG		(6 << 28)
+#define TPU_CLK_UART_DEBUG		(7 << 28)
+#define TPU_DEBUG_VALUE_MASK	((1 << 28) - 1)
+
+#define OSCCLK_RATE		24576000
+#define PLL_SHARED0_DIV0	1066000000
+#define PLL_SHARED1_DIV2	933000000
+#define PLL_SHARED2			800000000
+#define PLL_SHARED3			666000000
+#define PLL_SHARED0_DIV3	711000000
+#define PLL_SHARED1_DIV3	622000000
+#define PLL_SHARED0_DIV4	533000000
+#define PLL_SHARED2_DIV2	400000000
+#define PLL_SHARED3_DIV2	333000000
+
+#define TPU_CMU_TOP_REG 0x1E080000
+#define MUX_CLKCMU_TPU_TPU			(TPU_CMU_TPU_REG + 0x10FC)
+#define DIV_CLKCMU_TPU_TPU			(TPU_CMU_TPU_REG + 0x18EC)
+
+#define MUX_CLKCMU_TPU_TPUCTL		(TPU_CMU_TPU_REG + 0x1100)
+#define DIV_CLKCMU_TPU_TPUCTL		(TPU_CMU_TPU_REG + 0x18F0)
+
+#define MUX_CLKCMU_TPU_BUS			(TPU_CMU_TPU_REG + 0x10F8)
+#define DIV_CLKCMU_TPU_BUS			(TPU_CMU_TPU_REG + 0x18E8)
+
+#define MUX_CLKCMU_TPU_UART			(TPU_CMU_TPU_REG + 0x1104)
+#define DIV_CLKCMU_TPU_UART			(TPU_CMU_TPU_REG + 0x18F4)
+
+#define TPU_CMU_TPU_REG 0x1CC00000
+
+#define PLL_CON0_PLL_TPU			(TPU_CMU_TPU_REG + 0x0100)
+#define PLL_CON2_PLL_TPU			(TPU_CMU_TPU_REG + 0x0108)
+#define MUX_CLK_TPU_TPU				(TPU_CMU_TPU_REG + 0x1000)
+#define MUX_CLKCMU_TPU_TPU_USER		(TPU_CMU_TPU_REG + 0x0620)
+#define DIV_CLK_TPU_TPU				(TPU_CMU_TPU_REG + 0x1804)
+
+#define MUX_CLK_TPU_TPUCTL			(TPU_CMU_TPU_REG + 0x1004)
+#define MUX_CLKCMU_TPU_TPUCTL_USER	(TPU_CMU_TPU_REG + 0x0610)
+#define DIV_CLK_TPU_TPUCTL			(TPU_CMU_TPU_REG + 0x1808)
+
+#define MUX_CLKCMU_TPU_BUS_USER		(TPU_CMU_TPU_REG + 0x0600)
+
+#define DIV_CLK_TPU_BUSP			(TPU_CMU_TPU_REG + 0x1800)
+
+#define MUX_CLKCMU_TPU_UART_USER	(TPU_CMU_TPU_REG + 0x0630)
+
+#define MUX_USER_SEL_MASK		0x1
+#define MUX_USER_OSCCLK_SEL		0
+#define MUX_USER_CMUCLK_SEL		1
+
+#define MUX_CMU_SEL_MASK		0x7
+#define MUX_CMU_S0_D0_SEL		0
+#define MUX_CMU_S1_D2_SEL		1
+#define MUX_CMU_S2_SEL			2
+#define MUX_CMU_S3_SEL			3
+#define MUX_CMU_S0_D3_SEL		4
+#define MUX_CMU_S1_D3_SEL		5
+#define MUX_CMU_S0_D4_SEL		6
+
+#define MUX_CMU_UART_SEL_MASK		0x3
+#define MUX_CMU_UART_S0_D4_SEL		0
+#define MUX_CMU_UART_S2_D2_SEL		1
+#define MUX_CMU_UART_S3_D2_SEL		2
+
+#define DIV_CMU_RATIO_MASK		0xf
+#define DIV_USER_RATIO_MASK		0x7
+
+struct edgetpu_platform_pwr {
+	struct mutex policy_lock;
+	enum tpu_pwr_state curr_policy;
+};
 
 struct edgetpu_platform_dev {
 	struct edgetpu_dev edgetpu_dev;
+	struct edgetpu_platform_pwr platform_pwr;
 	int irq;
 	void *fw_region_vaddr;
 	size_t fw_region_size;
diff --git a/drivers/edgetpu/abrolhos-pwr.h b/drivers/edgetpu/abrolhos-pwr.h
index 40e29e6..e59858e 100644
--- a/drivers/edgetpu/abrolhos-pwr.h
+++ b/drivers/edgetpu/abrolhos-pwr.h
@@ -4,6 +4,35 @@
  *
  * Copyright (C) 2020 Google, Inc.
  */
+#ifndef __EDGETPU_PWR_H__
+#define __EDGETPU_PWR_H__
+
+/* Can't build out of tree with acpm_dvfs unless kernel supports ACPM */
+#if IS_ENABLED(CONFIG_ACPM_DVFS)
+
+#include <linux/acpm_dvfs.h>
+
+#else
+
+static int exynos_acpm_set_rate(unsigned int id, unsigned long rate)
+{
+	return 0;
+}
+static int exynos_acpm_set_init_freq(unsigned int dfs_id, unsigned long freq)
+{
+	return 0;
+}
+static unsigned long exynos_acpm_get_rate(unsigned int id,
+					  unsigned long dbg_val)
+{
+	return 0;
+}
+static int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
+{
+	return 0;
+}
+
+#endif /* IS_ENABLED(CONFIG_ACPM_DVFS) */
 
 /*
  * TPU Power States:
@@ -32,3 +61,7 @@ enum tpu_pwr_state {
 	TPU_ACTIVE_NOM = 1066000000,
 	TPU_ACTIVE_OD  = 1230000000,
 };
+
+#define TPU_POLICY_MAX	TPU_ACTIVE_OD
+
+#endif /* __EDGETPU_PWR_H__ */
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 248541d..0f94e67 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -21,6 +21,7 @@
 #include <linux/types.h>
 
 #include "edgetpu.h"
+#include "edgetpu-thermal.h"
 
 #define etdev_err(etdev, fmt, ...) dev_err((etdev)->etcdev, fmt, ##__VA_ARGS__)
 #define etdev_warn(etdev, fmt, ...)                                            \
@@ -143,6 +144,7 @@ struct edgetpu_dev {
 	struct edgetpu_kci *kci;
 	struct edgetpu_firmware *firmware; /* firmware management */
 	struct edgetpu_telemetry_ctx *telemetry;
+	struct edgetpu_thermal *thermal;
 	int mcp_id;		/* multichip pkg id, or -1 for none */
 	uint mcp_die_index;	/* physical die index w/in multichip pkg */
 	u8 mcp_pkg_type;	/* multichip pkg type */
diff --git a/drivers/edgetpu/edgetpu-thermal.c b/drivers/edgetpu/edgetpu-thermal.c
new file mode 100644
index 0000000..7561c5f
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-thermal.c
@@ -0,0 +1,261 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Edge TPU thermal driver.
+ *
+ * Copyright (C) 2020 Google, Inc.
+ */
+
+#include <linux/debugfs.h>
+#include <linux/device.h>
+#include <linux/gfp.h>
+#include <linux/mutex.h>
+#include <linux/pm_runtime.h>
+#include <linux/printk.h>
+#include <linux/slab.h>
+#include <linux/thermal.h>
+
+#include "abrolhos-firmware.h"
+#include "abrolhos-platform.h"
+#include "abrolhos-pwr.h"
+#include "edgetpu-config.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-mmu.h"
+#include "edgetpu-thermal.h"
+
+const static unsigned long state_mapping[] = {
+	TPU_ACTIVE_OD,
+	TPU_ACTIVE_NOM,
+	TPU_ACTIVE_UD,
+	TPU_ACTIVE_SUD,
+	TPU_RETENTION_CLOCKS_SLOW,
+	TPU_SLEEP_CLOCKS_SLOW,
+	TPU_SLEEP_CLOCKS_OFF,
+	TPU_DEEP_SLEEP_CLOCKS_FAST,
+	TPU_DEEP_SLEEP_CLOCKS_SLOW,
+	TPU_DEEP_SLEEP_CLOCKS_OFF,
+	TPU_OFF,
+};
+
+/*
+ * Sequence need to be kept to make power increasing
+ * to make sure we always return the highest state.
+ */
+const static struct edgetpu_state_pwr state_pwr_map[] = {
+	{ TPU_ACTIVE_OD, 198 },
+	{ TPU_ACTIVE_NOM, 165 },
+	{ TPU_ACTIVE_UD, 131 },
+	{ TPU_ACTIVE_SUD, 102 },
+	{ TPU_SLEEP_CLOCKS_SLOW, 66 },
+	{ TPU_SLEEP_CLOCKS_OFF, 66 },
+	{ TPU_RETENTION_CLOCKS_SLOW, 49 },
+	{ TPU_DEEP_SLEEP_CLOCKS_FAST, 43 },
+	{ TPU_DEEP_SLEEP_CLOCKS_SLOW, 6 },
+	{ TPU_DEEP_SLEEP_CLOCKS_OFF, 6 },
+	{ TPU_OFF, 0 },
+};
+
+static int edgetpu_get_max_state(struct thermal_cooling_device *cdev,
+				 unsigned long *state)
+{
+	*state = ARRAY_SIZE(state_mapping) - 1;
+	return 0;
+}
+
+/* Set cooling state
+ * Re-using code from abrohlos-platform.
+ * TODO: move to external call
+ */
+static int edgetpu_set_cur_state(struct thermal_cooling_device *cdev,
+				 unsigned long state_original)
+{
+	int ret;
+	struct edgetpu_thermal *cooling = cdev->devdata;
+	struct device *dev = cooling->dev;
+	unsigned long pwr_state;
+
+	if (state_original >= ARRAY_SIZE(state_mapping)) {
+		dev_err(dev, "%s: invalid cooling state %lu\n",
+			__func__, state_original);
+		WARN_ON(1);
+		return -EINVAL;
+	}
+
+	pwr_state = state_mapping[state_original];
+	cooling->pwr_state = pwr_state;
+
+	mutex_lock(&cooling->lock);
+
+	ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, pwr_state);
+	if (ret) {
+		dev_err(dev, "error setting tpu policy: %d\n", ret);
+		mutex_unlock(&cooling->lock);
+		return ret;
+	}
+
+	mutex_unlock(&cooling->lock);
+	return 0;
+}
+
+static int edgetpu_get_cur_state(struct thermal_cooling_device *cdev,
+				 unsigned long *state)
+{	unsigned long state_original;
+	struct edgetpu_thermal *cooling = cdev->devdata;
+	int i = 0;
+
+	state_original = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	find_state_pwr(i, state_original, state_mapping[i],
+		state_mapping, *state, i);
+	dev_err(cooling->dev, "Unknown get state req for: %lu\n",
+		state_original);
+	*state = 0;
+	WARN_ON(1);
+	return -EINVAL;
+}
+
+static int edgetpu_state2power_internal(unsigned long state, u32 *power,
+					struct device *dev)
+{
+	int i = 0;
+
+	find_state_pwr(i, state, state_pwr_map[i].state,
+		state_pwr_map, *power, state_pwr_map[i].power);
+	dev_err(dev, "Unknown state req for: %lu\n", state);
+	*power = 0;
+	WARN_ON(1);
+	return -EINVAL;
+}
+
+static int edgetpu_get_requested_power(struct thermal_cooling_device *cdev,
+				       struct thermal_zone_device *tz,
+				       u32 *power)
+{
+	int err = 0;
+	unsigned long state_original;
+	struct edgetpu_thermal *cooling = cdev->devdata;
+
+	state_original = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	err = edgetpu_state2power_internal(state_original, power,
+					   cooling->dev);
+	return err;
+}
+
+static int edgetpu_state2power(struct thermal_cooling_device *cdev,
+			       struct thermal_zone_device *tz,
+			       unsigned long state, u32 *power)
+{
+	int err = 0;
+	struct edgetpu_thermal *cooling = cdev->devdata;
+
+	if (state >= ARRAY_SIZE(state_mapping)) {
+		dev_err(cooling->dev, "%s: invalid state: %lu\n",
+		__func__, state);
+		return -EINVAL;
+	}
+
+	err = edgetpu_state2power_internal(state_mapping[state], power,
+					   cooling->dev);
+	return err;
+}
+
+static int edgetpu_power2state(struct thermal_cooling_device *cdev,
+			       struct thermal_zone_device *tz, u32 power,
+			       unsigned long *state)
+{
+	int i = 0;
+	struct edgetpu_thermal *cooling = cdev->devdata;
+
+	for (i = 0; i < ARRAY_SIZE(state_pwr_map); i++) {
+		if (power >= state_pwr_map[i].power) {
+			*state = state_pwr_map[i].state;
+			return 0;
+		}
+	}
+
+	dev_err(cooling->dev, "No power2state mapping found: %d\n",
+		power);
+	WARN_ON(1);
+	return -EINVAL;
+}
+
+static struct thermal_cooling_device_ops edgetpu_cooling_ops = {
+	.get_max_state = edgetpu_get_max_state,
+	.get_cur_state = edgetpu_get_cur_state,
+	.set_cur_state = edgetpu_set_cur_state,
+	.get_requested_power = edgetpu_get_requested_power,
+	.state2power = edgetpu_state2power,
+	.power2state = edgetpu_power2state,
+};
+
+static void tpu_thermal_exit_cooling(struct edgetpu_thermal *thermal)
+{
+	thermal_cooling_device_unregister(thermal->cdev);
+}
+
+static void tpu_thermal_exit(struct edgetpu_thermal *thermal)
+{
+	tpu_thermal_exit_cooling(thermal);
+	debugfs_remove_recursive(thermal->cooling_root);
+}
+
+static void devm_tpu_thermal_release(struct device *dev, void *res)
+{
+	struct edgetpu_thermal *thermal = res;
+
+	tpu_thermal_exit(thermal);
+}
+
+static int
+tpu_thermal_cooling_register(struct edgetpu_thermal *thermal, char *type)
+{
+	int err;
+
+	thermal->op_data = NULL;
+
+	mutex_init(&thermal->lock);
+	thermal->pwr_state = TPU_OFF;
+	thermal->cdev = thermal_of_cooling_device_register(
+		NULL, type, thermal, &edgetpu_cooling_ops);
+	if (IS_ERR(thermal->cdev)) {
+		err = PTR_ERR(thermal->cdev);
+		return err;
+	}
+	return 0;
+}
+
+static int tpu_thermal_init(struct edgetpu_thermal *thermal, struct device *dev)
+{
+	int err;
+
+	thermal->dev = dev;
+	thermal->cooling_root = debugfs_create_dir("edgetpu_cooling", NULL);
+
+	err = tpu_thermal_cooling_register(thermal,
+				       EDGETPU_COOLING_NAME);
+	if (err) {
+		dev_err(dev, "failed to initialize external cooling\n");
+		tpu_thermal_exit(thermal);
+		return err;
+	}
+
+	return 0;
+}
+
+struct edgetpu_thermal *devm_tpu_thermal_create(struct device *dev)
+{
+	struct edgetpu_thermal *thermal;
+	int err;
+
+	thermal = devres_alloc(devm_tpu_thermal_release,
+			       sizeof(struct edgetpu_thermal), GFP_KERNEL);
+	if (!thermal)
+		return ERR_PTR(-ENOMEM);
+
+	err = tpu_thermal_init(thermal, dev);
+	if (err) {
+		devres_free(thermal);
+		return ERR_PTR(err);
+	}
+
+	devres_add(dev, thermal);
+	return thermal;
+}
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
new file mode 100644
index 0000000..227b6ef
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -0,0 +1,48 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Edge TPU thermal driver header
+ *
+ * Copyright (C) 2020 Google, Inc.
+ */
+
+#ifndef __EDGETPU_THERMAL_H__
+#define __EDGETPU_THERMAL_H__
+
+#include <linux/debugfs.h>
+#include <linux/device.h>
+#include <linux/gfp.h>
+#include <linux/mutex.h>
+#include <linux/pm_runtime.h>
+#include <linux/printk.h>
+#include <linux/slab.h>
+#include <linux/thermal.h>
+
+#define find_state_pwr(i, cmp_left, cmp_right, list, out_left, out_right) \
+do { \
+		if (cmp_left == cmp_right) { \
+			out_left = out_right; \
+			return 0; \
+		} \
+		i++; \
+} while (i < ARRAY_SIZE(list))
+
+#define EDGETPU_COOLING_NAME "tpu_cooling"
+
+struct edgetpu_thermal {
+	struct device *dev;
+	struct dentry *cooling_root;
+	struct thermal_cooling_device *cdev;
+	struct mutex lock;
+	void *op_data;
+	unsigned long pwr_state;
+};
+
+struct edgetpu_state_pwr {
+	unsigned long state;
+	u32 power;
+};
+
+
+struct edgetpu_thermal *devm_tpu_thermal_create(struct device *dev);
+
+#endif /* __EDGETPU_THERMAL_H__ */