Merge branch 'pro' into android13-gs-pixel-5.10

  edgetpu: janeiro remove compat. code for shareability
  Bug: 193593081
  edgetpu: janeiro set shareability in after probe
  Bug: 200923214
  edgetpu: janeiro: Add TPU_OFF
  Bug: 199365799
  edgetpu: set first_open to false for external mailbox
  Bug: 200507456
  edgetpu: Put sscoredump features behind config flag
  Bug: 197672212
  edgetpu: rio import drivers
  Bug: 198718365
  Revert "edgetpu: remove edgetpu_unregister_irq"
  Bug: 200183185
  edgetpu: implement commands for interop
  Bug: 199355448
  edgetpu: fix up header inclusions
  edgetpu: janeiro use common mobile interface
  Bug: 199717700
  edgetpu: add edgetpu-mobile-platform.c
  Bug: 199717700
  edgetpu: janeiro use IRQ array in mobile dev
  Bug: 199717700
  edgetpu: remove edgetpu_unregister_irq
  edgetpu: mobile add IRQ fields to platform_dev
  Bug: 199717700
  edgetpu: support predefined KCI timeout
  Bug: 197068927
  edgetpu: mobile cpu reset according to #cores
  Bug: 199571462
  Bug: 198718365
  edgetpu: config add EDGETPU_NUM_CORES
  Bug: 199571462
  edgetpu: clearer config selection
  Bug: 198718365
  edgetpu: mobile firmware fix typos
  edgetpu: fix firmware load at different priv level
  edgetpu: fix returning/check correct status of errors
  Bug: 197919435
  edgetpu: janeiro: add GSA device
  Bug: 194054265
  edgetpu: janeiro: setup SSMT if GSA is available
  Bug: 197301774
  edgetpu: merge mobile firmware code
  Bug: 176881607
  edgetpu: merge mobile PM code
  Bug: 176881607
  edgetpu: move number of contexts to chip configs
  Bug: 194989925

GitOrigin-RevId: 96ee87151266f7bd4f96e5e3c4ae4cf1580bc8d2
Change-Id: If9c902d6d17b00502e01b0ba4f801985912640f6

30 files changed, 1856 insertions, 1130 deletions

diff --git a/drivers/edgetpu/BUILD.bazel b/drivers/edgetpu/BUILD.bazel
deleted file mode 100644
index 83d0c84..0000000
--- a/drivers/edgetpu/BUILD.bazel
+++ /dev/null
@@ -1,20 +0,0 @@
-# NOTE: THIS FILE IS EXPERIMENTAL FOR THE BAZEL MIGRATION AND NOT USED FOR
-# YOUR BUILDS CURRENTLY.
-#
-# It is not yet the source of truth for your build. If you're looking to modify
-# the build file, modify the Android.bp file instead. Do *not* modify this file
-# unless you have coordinated with the team managing the Soong to Bazel
-# migration.
-
-load("//build/kleaf:kernel.bzl", "kernel_module")
-
-kernel_module(
-    name = "edgetpu.cloudripper",
-    outs = [
-        "janeiro.ko",
-    ],
-    kernel_build = "//private/gs-google:cloudripper",
-    visibility = [
-        "//private/gs-google:__pkg__",
-    ],
-)
diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 88d8163..3479ad2 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -12,10 +12,9 @@ endif
 
 edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o edgetpu-firmware.o edgetpu-firmware-util.o
 
-edgetpu-mobile-objs	:= edgetpu-mobile-firmware.o
 
+janeiro-y	:= janeiro-device.o janeiro-device-group.o janeiro-fs.o janeiro-core.o janeiro-platform.o janeiro-firmware.o janeiro-pm.o janeiro-debug-dump.o janeiro-usage-stats.o janeiro-iommu.o janeiro-wakelock.o janeiro-external.o $(edgetpu-objs)
 
-janeiro-y	:= janeiro-device.o janeiro-device-group.o janeiro-fs.o janeiro-core.o janeiro-platform.o janeiro-firmware.o janeiro-pm.o janeiro-debug-dump.o janeiro-usage-stats.o janeiro-iommu.o janeiro-wakelock.o $(edgetpu-mobile-objs) $(edgetpu-objs)
 
 CFLAGS_janeiro-fs.o := -DCONFIG_JANEIRO=1
 CFLAGS_janeiro-core.o := -DCONFIG_JANEIRO=1
@@ -29,3 +28,4 @@ CFLAGS_janeiro-thermal.o := -DCONFIG_JANEIRO=1
 CFLAGS_janeiro-debug-dump.o := -DCONFIG_JANEIRO=1
 CFLAGS_janeiro-usage-stats.o := -DCONFIG_JANEIRO=1
 CFLAGS_janeiro-wakelock.o := -DCONFIG_JANEIRO=1
+CFLAGS_janeiro-external.o := -DCONFIG_JANEIRO=1
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index 3bbfaaa..f21f1dc 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -18,14 +18,12 @@ edgetpu-objs	:= edgetpu-async.o edgetpu-dmabuf.o edgetpu-iremap-pool.o \
 		   edgetpu-sw-watchdog.o edgetpu-telemetry.o \
 		   edgetpu-firmware-util.o edgetpu-firmware.o
 
-edgetpu-mobile-objs	:= edgetpu-mobile-firmware.o
-
-
 janeiro-objs	:= janeiro-core.o janeiro-debug-dump.o janeiro-device-group.o \
 		   janeiro-device.o janeiro-firmware.o janeiro-fs.o \
 		   janeiro-iommu.o janeiro-platform.o janeiro-pm.o \
-		   janeiro-usage-stats.o janeiro-wakelock.o \
-		   $(edgetpu-objs) $(edgetpu-mobile-objs)
+		   janeiro-usage-stats.o janeiro-wakelock.o janeiro-external.o \
+		   $(edgetpu-objs)
+
 
 KBUILD_OPTIONS += CONFIG_JANEIRO=m
 
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index 9306fac..0626ba0 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -8,11 +8,20 @@
 #ifndef __EDGETPU_CONFIG_H__
 #define __EDGETPU_CONFIG_H__
 
-#ifdef CONFIG_JANEIRO
-
+#if IS_ENABLED(CONFIG_JANEIRO)
 #include "janeiro/config.h"
-#endif /* CONFIG_JANEIRO */
+
+#else /* unknown */
+
+#error "Unknown EdgeTPU config"
+
+#endif /* unknown */
+
 #define EDGETPU_DEFAULT_FIRMWARE_NAME "google/edgetpu-" DRIVER_NAME ".fw"
 #define EDGETPU_TEST_FIRMWARE_NAME "google/edgetpu-" DRIVER_NAME "-test.fw"
 
+#ifndef EDGETPU_NUM_CORES
+#define EDGETPU_NUM_CORES 1
+#endif
+
 #endif /* __EDGETPU_CONFIG_H__ */
diff --git a/drivers/edgetpu/edgetpu-external.c b/drivers/edgetpu/edgetpu-external.c
new file mode 100644
index 0000000..bbbea8c
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-external.c
@@ -0,0 +1,134 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Utility functions for interfacing other modules with Edge TPU ML accelerator.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#include <linux/device.h>
+#include <linux/err.h>
+
+#include <soc/google/tpu-ext.h>
+
+#include "edgetpu-config.h"
+#include "edgetpu-device-group.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-mailbox.h"
+
+static int edgetpu_external_mailbox_info_get(struct edgetpu_external_mailbox_info *info,
+					     struct edgetpu_external_mailbox *ext_mailbox)
+{
+	int i;
+	u32 count = ext_mailbox->count;
+	struct edgetpu_mailbox_descriptor *desc;
+
+	if (!info)
+		return -EINVAL;
+
+	if (info->count < count) {
+		etdev_err(ext_mailbox->etdev,
+			  "Insufficient space in provided buffer expected: %d received: %d\n",
+			  count, info->count);
+		return -ENOMEM;
+	}
+
+	for (i = 0; i < count; i++) {
+		desc = &ext_mailbox->descriptors[i];
+		info->mailboxes[i].cmdq_pa = desc->cmd_queue_mem.phys_addr;
+		info->mailboxes[i].respq_pa = desc->resp_queue_mem.phys_addr;
+		info->mailboxes[i].mailbox_id =
+			desc->mailbox->mailbox_id - (EDGETPU_NUM_VII_MAILBOXES + 1);
+	}
+
+	info->cmdq_size = ext_mailbox->attr.cmd_queue_size;
+	info->respq_size = ext_mailbox->attr.resp_queue_size;
+	info->count = count;
+
+	return 0;
+}
+
+static bool is_edgetpu_valid_client(struct edgetpu_external_mailbox *ext_mailbox,
+				    enum edgetpu_external_client_type client_type)
+{
+	switch (client_type) {
+	case EDGETPU_EXTERNAL_CLIENT_TYPE_DSP:
+		return ext_mailbox->client_type == EDGETPU_EXT_MAILBOX_TYPE_DSP;
+	default:
+		return false;
+	}
+
+	return false;
+}
+
+static int edgetpu_mailbox_external_info_get_cmd(struct device *edgetpu_dev,
+						 enum edgetpu_external_client_type client_type,
+						 int *client_fd,
+						 struct edgetpu_external_mailbox_info *info)
+{
+	struct edgetpu_client *client;
+	struct edgetpu_device_group *group;
+	struct edgetpu_external_mailbox *ext_mailbox;
+	int ret = 0;
+	struct fd f = fdget(*client_fd);
+	struct file *file = f.file;
+
+	if (!file)
+		return -ENOENT;
+
+	if (!is_edgetpu_file(file)) {
+		ret = -ENOENT;
+		goto out;
+	}
+
+	client = file->private_data;
+	if (!client) {
+		ret = -EINVAL;
+		goto out;
+	}
+
+	mutex_lock(&client->group_lock);
+	if (!client->group || client->etdev->dev != edgetpu_dev) {
+		ret = -EINVAL;
+		mutex_unlock(&client->group_lock);
+		goto out;
+	}
+	group = edgetpu_device_group_get(client->group);
+	mutex_unlock(&client->group_lock);
+
+	mutex_lock(&group->lock);
+	ext_mailbox = group->ext_mailbox;
+	if (!ext_mailbox) {
+		ret = -ENOENT;
+		goto unlock;
+	}
+	if (!is_edgetpu_valid_client(ext_mailbox, client_type)) {
+		ret = -EINVAL;
+		goto unlock;
+	}
+	ret = edgetpu_external_mailbox_info_get(info, ext_mailbox);
+unlock:
+	mutex_unlock(&group->lock);
+	edgetpu_device_group_put(group);
+out:
+	fdput(f);
+	return ret;
+}
+
+int edgetpu_ext_driver_cmd(struct device *edgetpu_dev,
+			   enum edgetpu_external_client_type client_type,
+			   enum edgetpu_external_commands cmd_id, void *in_data, void *out_data)
+{
+	int ret = 0;
+
+	switch (cmd_id) {
+	case MAILBOX_EXTERNAL_INFO_GET:
+		ret = edgetpu_mailbox_external_info_get_cmd(edgetpu_dev, client_type, in_data,
+							    out_data);
+		break;
+	default:
+		return -ENOENT;
+	}
+
+	return ret;
+}
+EXPORT_SYMBOL_GPL(edgetpu_ext_driver_cmd);
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 8b5a782..46de18f 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -66,7 +66,7 @@ static struct dentry *edgetpu_debugfs_dir;
 		}                                                              \
 	} while (0)
 
-static bool is_edgetpu_file(struct file *file)
+bool is_edgetpu_file(struct file *file)
 {
 	if (edgetpu_is_external_wrapper_class_file(file))
 		return true;
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index f15b61a..eae1730 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -14,6 +14,7 @@
 #include <linux/types.h>
 #include <linux/version.h>
 
+#include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-mapping.h"
 #include "edgetpu-mmu.h"
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index f223dcb..6642e5f 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -58,12 +58,6 @@
 /* The number of TPU tiles in an edgetpu chip */
 #define EDGETPU_NTILES	16
 
-/* Up to 7 concurrent device groups / workloads per device. */
-#define EDGETPU_NGROUPS		7
-
-/* 1 context per VII/group plus 1 for KCI */
-#define EDGETPU_NCONTEXTS	(EDGETPU_NGROUPS + 1)
-
 /*
  * Common-layer context IDs for non-secure TPU access, translated to chip-
  * specific values in the mmu driver.
@@ -71,8 +65,7 @@
 enum edgetpu_context_id {
 	EDGETPU_CONTEXT_INVALID = -1,
 	EDGETPU_CONTEXT_KCI = 0,	/* TPU firmware/kernel ID 0 */
-	EDGETPU_CONTEXT_VII_BASE = 1,	/* groups 0-6 IDs 1-7 */
-	/* contexts 8 and above not yet allocated */
+	EDGETPU_CONTEXT_VII_BASE = 1,	/* groups IDs starts from 1 to (EDGETPU_CONTEXTS - 1) */
 	/* A bit mask to mark the context is an IOMMU domain token */
 	EDGETPU_CONTEXT_DOMAIN_TOKEN = 1 << 30,
 };
@@ -232,8 +225,6 @@ struct edgetpu_dev {
 	/* debug dump handlers */
 	edgetpu_debug_dump_handlers *debug_dump_handlers;
 	struct work_struct debug_dump_work;
-	/* status about if device going to be removed from system */
-	bool on_exit;
 };
 
 struct edgetpu_dev_iface {
@@ -358,6 +349,8 @@ void edgetpu_free_coherent(struct edgetpu_dev *etdev,
 			   struct edgetpu_coherent_mem *mem,
 			   enum edgetpu_context_id context_id);
 
+/* Checks if @file belongs to edgetpu driver */
+bool is_edgetpu_file(struct file *file);
 
 /* External drivers can hook up to edgetpu driver using these calls. */
 int edgetpu_open(struct edgetpu_dev_iface *etiface, struct file *file);
@@ -388,7 +381,9 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 		       const struct edgetpu_iface_params *iface_params,
 		       uint num_ifaces);
 void edgetpu_device_remove(struct edgetpu_dev *etdev);
+/* Registers IRQ. */
 int edgetpu_register_irq(struct edgetpu_dev *etdev, int irq);
+/* Reverts edgetpu_register_irq */
 void edgetpu_unregister_irq(struct edgetpu_dev *etdev, int irq);
 
 /* Core -> Device FS API */
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index d262bba..17d92ab 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -30,7 +30,11 @@
 #define QUEUE_SIZE MAX_QUEUE_SIZE
 
 /* Timeout for KCI responses from the firmware (milliseconds) */
-#if IS_ENABLED(CONFIG_EDGETPU_TEST)
+#ifdef EDGETPU_KCI_TIMEOUT
+
+#define KCI_TIMEOUT EDGETPU_KCI_TIMEOUT
+
+#elif IS_ENABLED(CONFIG_EDGETPU_TEST)
 /* fake-firmware could respond in a short time */
 #define KCI_TIMEOUT	(200)
 #else
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 56acd72..f4a3e4e 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -853,6 +853,10 @@ static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 	if (!ext_mailbox_req)
 		return -EINVAL;
 
+	ret = edgetpu_mailbox_validate_attr(&ext_mailbox_req->attr);
+	if (ret)
+		return ret;
+
 	count = ext_mailbox_req->count;
 	attr = ext_mailbox_req->attr;
 
@@ -874,6 +878,7 @@ static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 	ext_mailbox->attr = attr;
 	ext_mailbox->count = count;
 	ext_mailbox->etdev = group->etdev;
+	ext_mailbox->client_type = ext_mailbox_req->client_type;
 
 	write_lock_irqsave(&mgr->mailboxes_lock, flags);
 	for (i = ext_mailbox_req->start; i <= ext_mailbox_req->end; i++) {
@@ -892,6 +897,7 @@ static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 				mgr->mailboxes[i] = mailbox;
 				ext_mailbox->descriptors[j++].mailbox = mailbox;
 			} else {
+				ret = PTR_ERR(mailbox);
 				goto release;
 			}
 		}
@@ -989,7 +995,7 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 	for (i = 0; i < ext_mailbox->count; i++) {
 		id = ext_mailbox->descriptors[i].mailbox->mailbox_id;
 		etdev_dbg(group->etdev, "Enabling mailbox: %d\n", id);
-		ret = edgetpu_mailbox_activate(group->etdev, id, vcid, true);
+		ret = edgetpu_mailbox_activate(group->etdev, id, vcid, false);
 		if (ret) {
 			etdev_err(group->etdev, "Activate mailbox %d failed: %d", id, ret);
 			break;
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index d76534a..3098760 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -91,7 +91,12 @@ struct edgetpu_mailbox_descriptor {
 /* Structure to hold multiple external mailboxes allocated for a device group. */
 struct edgetpu_external_mailbox {
 	/* Number of external mailboxes allocated for a device group. */
-	int count;
+	u32 count;
+	/*
+	 * Client type of external mailboxes requester, it belongs to
+	 * EDGETPU_EXT_MAILBOX_TYPE_*
+	 */
+	u64 client_type;
 	/* Leader of device group. */
 	struct edgetpu_dev *etdev;
 	/* Array of external mailboxes info with length @count. */
@@ -106,6 +111,11 @@ struct edgetpu_external_mailbox_req {
 	uint end; /* end index of external mailbox in mailbox_manager */
 	/* number of mailboxes to be allocated, should be less or equal to (end - start + 1) */
 	uint count;
+	/*
+	 * Client type of external mailboxes requester, it belongs to
+	 * EDGETPU_EXT_MAILBOX_TYPE_*
+	 */
+	u64 client_type;
 	struct edgetpu_mailbox_attr attr; /* mailbox attribute for allocation */
 };
 
@@ -487,12 +497,7 @@ static inline void edgetpu_mailbox_enable(struct edgetpu_mailbox *mailbox)
 /* Disables a mailbox by setting CSR. */
 static inline void edgetpu_mailbox_disable(struct edgetpu_mailbox *mailbox)
 {
-	/*
-	 * if device is being removed, no need to access device CSR since it
-	 * is going to be powered-off anyway
-	 */
-	if (!mailbox->etdev->on_exit)
-		EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 0);
+	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 0);
 }
 
 #endif /* __EDGETPU_MAILBOX_H__ */
diff --git a/drivers/edgetpu/edgetpu-mobile-firmware.c b/drivers/edgetpu/edgetpu-mobile-firmware.c
deleted file mode 100644
index dedac0f..0000000
--- a/drivers/edgetpu/edgetpu-mobile-firmware.c
+++ /dev/null
@@ -1,63 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-/*
- * Edge TPU firmware management for mobile chipsets.
- *
- * Copyright (C) 2021 Google, Inc.
- */
-
-#include <linux/device.h>
-#include <linux/firmware.h>
-#include <linux/slab.h>
-#include <linux/string.h>
-#include <linux/types.h>
-
-#include "edgetpu-firmware.h"
-#include "edgetpu-firmware-util.h"
-#include "edgetpu-internal.h"
-#include "mobile-firmware.h"
-
-/* Load firmware for chips that use carveout memory for a single chip. */
-int edgetpu_firmware_chip_load_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc, const char *name)
-{
-	int ret;
-	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct device *dev = etdev->dev;
-	const struct firmware *fw;
-	size_t aligned_size;
-
-	ret = request_firmware(&fw, name, dev);
-	if (ret) {
-		etdev_dbg(etdev,
-			  "%s: request '%s' failed: %d\n", __func__, name, ret);
-		return ret;
-	}
-
-	aligned_size = ALIGN(fw->size, fw_desc->buf.used_size_align);
-	if (aligned_size > fw_desc->buf.alloc_size) {
-		etdev_dbg(etdev,
-			  "%s: firmware buffer too small: alloc size=%#zx, required size=%#zx\n",
-			  __func__, fw_desc->buf.alloc_size, aligned_size);
-		ret = -ENOSPC;
-		goto out_release_firmware;
-	}
-
-	memcpy(fw_desc->buf.vaddr, fw->data, fw->size);
-	fw_desc->buf.used_size = aligned_size;
-	/* May return NULL on out of memory, driver must handle properly */
-	fw_desc->buf.name = kstrdup(name, GFP_KERNEL);
-
-out_release_firmware:
-	release_firmware(fw);
-	return ret;
-}
-
-void edgetpu_firmware_chip_unload_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc)
-{
-	kfree(fw_desc->buf.name);
-	fw_desc->buf.name = NULL;
-	fw_desc->buf.used_size = 0;
-}
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
new file mode 100644
index 0000000..7b9abec
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -0,0 +1,340 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Common platform interfaces for mobile TPU chips.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#include <linux/device.h>
+#include <linux/gsa/gsa_tpu.h>
+#include <linux/module.h>
+#include <linux/of.h>
+#include <linux/of_address.h>
+#include <linux/of_platform.h>
+#include <linux/platform_device.h>
+
+#include "edgetpu-config.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-iremap-pool.h"
+#include "edgetpu-mmu.h"
+#include "edgetpu-mobile-platform.h"
+#include "edgetpu-telemetry.h"
+#include "mobile-firmware.h"
+#include "mobile-pm.h"
+
+/*
+ * Log and trace buffers at the beginning of the remapped region,
+ * pool memory afterwards.
+ */
+#define EDGETPU_POOL_MEM_OFFSET (EDGETPU_TELEMETRY_BUFFER_SIZE * 2)
+
+static void get_telemetry_mem(struct edgetpu_mobile_platform_dev *etmdev,
+			      enum edgetpu_telemetry_type type, struct edgetpu_coherent_mem *mem)
+{
+	int offset = type == EDGETPU_TELEMETRY_TRACE ? EDGETPU_TELEMETRY_BUFFER_SIZE : 0;
+
+	mem->vaddr = etmdev->shared_mem_vaddr + offset;
+	mem->dma_addr = EDGETPU_REMAPPED_DATA_ADDR + offset;
+	mem->tpu_addr = EDGETPU_REMAPPED_DATA_ADDR + offset;
+	mem->host_addr = 0;
+	mem->size = EDGETPU_TELEMETRY_BUFFER_SIZE;
+}
+
+static void edgetpu_mobile_get_telemetry_mem(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	get_telemetry_mem(etmdev, EDGETPU_TELEMETRY_LOG, &etmdev->log_mem);
+	get_telemetry_mem(etmdev, EDGETPU_TELEMETRY_TRACE, &etmdev->trace_mem);
+}
+
+static int edgetpu_platform_setup_fw_region(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct platform_device *gsa_pdev;
+	struct device *dev = etdev->dev;
+	struct resource r;
+	struct device_node *np;
+	int err;
+	size_t region_map_size = EDGETPU_FW_SIZE_MAX + EDGETPU_REMAPPED_DATA_SIZE;
+
+	np = of_parse_phandle(dev->of_node, "memory-region", 0);
+	if (!np) {
+		dev_err(dev, "No memory region for firmware");
+		return -ENODEV;
+	}
+
+	err = of_address_to_resource(np, 0, &r);
+	of_node_put(np);
+	if (err) {
+		dev_err(dev, "No memory address assigned to firmware region");
+		return err;
+	}
+
+	if (resource_size(&r) < region_map_size) {
+		dev_err(dev, "Memory region for firmware too small (%zu bytes needed, got %llu)",
+			region_map_size, resource_size(&r));
+		return -ENOSPC;
+	}
+
+	/* Get GSA device from device tree */
+	np = of_parse_phandle(dev->of_node, "gsa-device", 0);
+	if (!np) {
+		dev_warn(dev, "No gsa-device in device tree. Authentication not available");
+	} else {
+		gsa_pdev = of_find_device_by_node(np);
+		if (!gsa_pdev) {
+			dev_err(dev, "GSA device not found");
+			of_node_put(np);
+			return -ENODEV;
+		}
+		etmdev->gsa_dev = get_device(&gsa_pdev->dev);
+		of_node_put(np);
+	}
+
+	etmdev->fw_region_paddr = r.start;
+	etmdev->fw_region_size = EDGETPU_FW_SIZE_MAX;
+
+	etmdev->shared_mem_vaddr = memremap(r.start + EDGETPU_REMAPPED_DATA_OFFSET,
+					    EDGETPU_REMAPPED_DATA_SIZE, MEMREMAP_WC);
+	if (!etmdev->shared_mem_vaddr) {
+		dev_err(dev, "Shared memory remap failed");
+		if (etmdev->gsa_dev)
+			put_device(etmdev->gsa_dev);
+		return -EINVAL;
+	}
+	etmdev->shared_mem_paddr = r.start + EDGETPU_REMAPPED_DATA_OFFSET;
+
+	return 0;
+}
+
+static void edgetpu_platform_cleanup_fw_region(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	if (etmdev->gsa_dev) {
+		gsa_unload_tpu_fw_image(etmdev->gsa_dev);
+		put_device(etmdev->gsa_dev);
+	}
+	if (!etmdev->shared_mem_vaddr)
+		return;
+	memunmap(etmdev->shared_mem_vaddr);
+	etmdev->shared_mem_vaddr = NULL;
+}
+
+int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev)
+{
+	int ret;
+
+	ret = edgetpu_mmu_attach(etdev, NULL);
+	if (ret)
+		dev_err(etdev->dev, "failed to attach IOMMU: %d", ret);
+	return ret;
+}
+
+void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
+{
+	edgetpu_mmu_detach(etdev);
+}
+
+static int edgetpu_platform_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct platform_device *pdev = to_platform_device(etdev->dev);
+	struct resource *res;
+	int ret;
+	void __iomem *ssmt_base;
+
+	res = platform_get_resource_byname(pdev, IORESOURCE_MEM, "ssmt");
+	if (!res) {
+		etdev_warn(etdev, "Failed to find SSMT register base");
+		return -EINVAL;
+	}
+	ssmt_base = devm_ioremap_resource(&pdev->dev, res);
+	if (IS_ERR(ssmt_base)) {
+		ret = PTR_ERR(ssmt_base);
+		etdev_warn(etdev, "Failed to map SSMT register base: %d", ret);
+		return ret;
+	}
+	etmdev->ssmt_base = ssmt_base;
+	return 0;
+}
+
+static int edgetpu_platform_setup_irq(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct platform_device *pdev = to_platform_device(etdev->dev);
+	int n = platform_irq_count(pdev);
+	int ret;
+	int i;
+
+	etmdev->irq = devm_kmalloc_array(etdev->dev, n, sizeof(*etmdev->irq), GFP_KERNEL);
+	if (!etmdev->irq)
+		return -ENOMEM;
+
+	for (i = 0; i < n; i++) {
+		etmdev->irq[i] = platform_get_irq(pdev, i);
+		ret = edgetpu_register_irq(etdev, etmdev->irq[i]);
+		if (ret)
+			goto rollback;
+	}
+	etmdev->n_irq = n;
+	return 0;
+
+rollback:
+	while (i--)
+		edgetpu_unregister_irq(etdev, etmdev->irq[i]);
+	return ret;
+}
+
+static void edgetpu_platform_remove_irq(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	int i;
+
+	for (i = 0; i < etmdev->n_irq; i++)
+		edgetpu_unregister_irq(etdev, etmdev->irq[i]);
+}
+
+static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
+					 struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct device *dev = &pdev->dev;
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct resource *r;
+	struct edgetpu_mapped_resource regs;
+	int ret;
+	struct edgetpu_iface_params iface_params[] = {
+		/* Default interface  */
+		{ .name = NULL },
+		/* Common name for embedded SoC devices */
+		{ .name = "edgetpu-soc" },
+	};
+
+	mutex_init(&etmdev->tz_mailbox_lock);
+
+	platform_set_drvdata(pdev, etdev);
+	etdev->dev = dev;
+
+	r = platform_get_resource(pdev, IORESOURCE_MEM, 0);
+	if (IS_ERR_OR_NULL(r)) {
+		dev_err(dev, "failed to get memory resource");
+		return -ENODEV;
+	}
+
+	regs.phys = r->start;
+	regs.size = resource_size(r);
+	regs.mem = devm_ioremap_resource(dev, r);
+	if (IS_ERR_OR_NULL(regs.mem)) {
+		dev_err(dev, "failed to map registers");
+		return -ENODEV;
+	}
+
+	mutex_init(&etmdev->platform_pwr.policy_lock);
+	etmdev->platform_pwr.curr_policy = TPU_POLICY_MAX;
+
+	ret = edgetpu_chip_pm_create(etdev);
+	if (ret) {
+		dev_err(dev, "Failed to initialize PM interface: %d", ret);
+		return ret;
+	}
+
+	ret = edgetpu_platform_setup_fw_region(etmdev);
+	if (ret) {
+		dev_err(dev, "setup fw regions failed: %d", ret);
+		goto out_shutdown;
+	}
+
+	ret = edgetpu_iremap_pool_create(etdev,
+					 /* Base virtual address (kernel address space) */
+					 etmdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
+					 /* Base DMA address */
+					 EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
+					 /* Base TPU address */
+					 EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
+					 /* Base physical address */
+					 etmdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
+					 /* Size */
+					 EDGETPU_REMAPPED_DATA_SIZE - EDGETPU_POOL_MEM_OFFSET,
+					 /* Granularity */
+					 PAGE_SIZE);
+	if (ret) {
+		dev_err(dev, "failed to initialize remapped memory pool: %d", ret);
+		goto out_cleanup_fw;
+	}
+
+	etdev->mcp_id = -1;
+	etdev->mcp_die_index = 0;
+	ret = edgetpu_device_add(etdev, &regs, iface_params, ARRAY_SIZE(iface_params));
+	if (ret) {
+		dev_err(dev, "edgetpu setup failed: %d", ret);
+		goto out_destroy_iremap;
+	}
+
+	ret = edgetpu_platform_setup_irq(etmdev);
+	if (ret) {
+		dev_err(dev, "IRQ setup failed: %d", ret);
+		goto out_remove_device;
+	}
+
+	ret = edgetpu_platform_parse_ssmt(etmdev);
+	if (ret)
+		dev_warn(dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
+
+	edgetpu_mobile_get_telemetry_mem(etmdev);
+	ret = edgetpu_telemetry_init(etdev, &etmdev->log_mem, &etmdev->trace_mem);
+	if (ret)
+		goto out_remove_irq;
+
+	ret = edgetpu_mobile_firmware_create(etdev);
+	if (ret) {
+		dev_err(dev, "initialize firmware downloader failed: %d", ret);
+		goto out_tel_exit;
+	}
+
+	if (etmdev->after_probe) {
+		ret = etmdev->after_probe(etmdev);
+		if (ret) {
+			dev_err(dev, "after_probe callback failed: %d", ret);
+			goto out_destroy_fw;
+		}
+	}
+
+	dev_info(dev, "%s edgetpu initialized. Build: %s", etdev->dev_name, GIT_REPO_TAG);
+	/* Turn the device off unless a client request is already received. */
+	edgetpu_pm_shutdown(etdev, false);
+
+	return 0;
+out_destroy_fw:
+	edgetpu_mobile_firmware_destroy(etdev);
+out_tel_exit:
+	edgetpu_telemetry_exit(etdev);
+out_remove_irq:
+	edgetpu_platform_remove_irq(etmdev);
+out_remove_device:
+	edgetpu_device_remove(etdev);
+out_destroy_iremap:
+	edgetpu_iremap_pool_destroy(etdev);
+out_cleanup_fw:
+	edgetpu_platform_cleanup_fw_region(etmdev);
+out_shutdown:
+	dev_dbg(dev, "Probe finished with error %d, powering down", ret);
+	edgetpu_pm_shutdown(etdev, true);
+	return ret;
+}
+
+static int edgetpu_mobile_platform_remove(struct platform_device *pdev)
+{
+	struct edgetpu_dev *etdev = platform_get_drvdata(pdev);
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+
+	if (etmdev->before_remove)
+		etmdev->before_remove(etmdev);
+	edgetpu_mobile_firmware_destroy(etdev);
+	edgetpu_platform_remove_irq(etmdev);
+	edgetpu_pm_get(etdev->pm);
+	edgetpu_telemetry_exit(etdev);
+	edgetpu_device_remove(etdev);
+	edgetpu_iremap_pool_destroy(etdev);
+	edgetpu_platform_cleanup_fw_region(etmdev);
+	edgetpu_pm_put(etdev->pm);
+	edgetpu_pm_shutdown(etdev, true);
+	mobile_pm_destroy(etdev);
+	return 0;
+}
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
new file mode 100644
index 0000000..8b3a25f
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -0,0 +1,120 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Common platform interfaces for mobile TPU chips.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#ifndef __EDGETPU_MOBILE_PLATFORM_H__
+#define __EDGETPU_MOBILE_PLATFORM_H__
+
+#include <linux/debugfs.h>
+#include <linux/device.h>
+#include <linux/gfp.h>
+#include <linux/kernel.h>
+#include <linux/mutex.h>
+#include <linux/types.h>
+#include <soc/google/exynos_pm_qos.h>
+
+#if IS_ENABLED(CONFIG_GOOGLE_BCL)
+#include <soc/google/bcl.h>
+#endif
+
+#include "edgetpu-config.h"
+#include "edgetpu-internal.h"
+
+#define to_mobile_dev(etdev) container_of(etdev, struct edgetpu_mobile_platform_dev, edgetpu_dev)
+
+struct edgetpu_mobile_platform_pwr {
+	struct dentry *debugfs_dir;
+	struct mutex policy_lock;
+	enum edgetpu_pwr_state curr_policy;
+	struct mutex state_lock;
+	u64 min_state;
+	u64 requested_state;
+	/* INT/MIF requests for memory bandwidth */
+	struct exynos_pm_qos_request int_min;
+	struct exynos_pm_qos_request mif_min;
+	/* BTS */
+	unsigned int performance_scenario;
+	int scenario_count;
+	struct mutex scenario_lock;
+
+	/* LPM callbacks, NULL for chips without LPM */
+	int (*lpm_up)(struct edgetpu_dev *etdev);
+	void (*lpm_down)(struct edgetpu_dev *etdev);
+
+	/* Block shutdown callback, may be NULL */
+	void (*block_down)(struct edgetpu_dev *etdev);
+
+	/* Firmware shutdown callback. Must be implemented */
+	void (*firmware_down)(struct edgetpu_dev *etdev);
+
+	/* Chip-specific setup after the PM interface is created */
+	int (*after_create)(struct edgetpu_dev *etdev);
+
+	/* Chip-specific cleanup before the PM interface is destroyed */
+	int (*before_destroy)(struct edgetpu_dev *etdev);
+
+	/* ACPM set rate callback. Must be implemented */
+	int (*acpm_set_rate)(unsigned int id, unsigned long rate);
+};
+
+struct edgetpu_mobile_platform_dev {
+	/* Generic edgetpu device */
+	struct edgetpu_dev edgetpu_dev;
+	/* Common mobile platform power interface */
+	struct edgetpu_mobile_platform_pwr platform_pwr;
+	/* Physical address of the firmware image */
+	phys_addr_t fw_region_paddr;
+	/* Size of the firmware region */
+	size_t fw_region_size;
+	/* Virtual address of the memory region shared with firmware */
+	void *shared_mem_vaddr;
+	/* Physical address of the memory region shared with firmware */
+	phys_addr_t shared_mem_paddr;
+	/* Size of the shared memory region size */
+	size_t shared_mem_size;
+	/*
+	 * Pointer to GSA device for firmware authentication.
+	 * May be NULL if the chip does not support firmware authentication
+	 */
+	struct device *gsa_dev;
+	/* Virtual address of the SSMT block for this chip. */
+	void __iomem *ssmt_base;
+	/* Coherent log buffer */
+	struct edgetpu_coherent_mem log_mem;
+	/* Coherent trace buffer */
+	struct edgetpu_coherent_mem trace_mem;
+#if IS_ENABLED(CONFIG_GOOGLE_BCL)
+	struct bcl_device *bcl_dev;
+#endif
+	/* Protects TZ Mailbox client pointer */
+	struct mutex tz_mailbox_lock;
+	/* TZ mailbox client */
+	struct edgetpu_client *secure_client;
+
+	/* Length of @irq */
+	int n_irq;
+	/* Array of IRQ numbers */
+	int *irq;
+
+	/* callbacks for chip-dependent implementations */
+
+	/*
+	 * Called when common device probing procedure is done.
+	 *
+	 * Return a non-zero value can fail the probe procedure.
+	 *
+	 * This callback is optional.
+	 */
+	int (*after_probe)(struct edgetpu_mobile_platform_dev *etmdev);
+	/*
+	 * Called before common device removal procedure.
+	 *
+	 * This callback is optional.
+	 */
+	void (*before_remove)(struct edgetpu_mobile_platform_dev *etmdev);
+};
+
+#endif /* __EDGETPU_MOBILE_PLATFORM_H__ */
diff --git a/drivers/edgetpu/janeiro-debug-dump.c b/drivers/edgetpu/janeiro-debug-dump.c
index 394ec08..4314abe 100644
--- a/drivers/edgetpu/janeiro-debug-dump.c
+++ b/drivers/edgetpu/janeiro-debug-dump.c
@@ -5,7 +5,7 @@
  * Copyright (C) 2021 Google, Inc.
  */
 
-#include "mobile-debug-dump.c"
+#include "edgetpu-debug-dump.c"
 
 int edgetpu_debug_dump_init(struct edgetpu_dev *etdev)
 {
diff --git a/drivers/edgetpu/janeiro-device.c b/drivers/edgetpu/janeiro-device.c
index 36dbe76..cc4c975 100644
--- a/drivers/edgetpu/janeiro-device.c
+++ b/drivers/edgetpu/janeiro-device.c
@@ -11,8 +11,8 @@
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-mailbox.h"
+#include "edgetpu-mobile-platform.h"
 #include "edgetpu-telemetry.h"
-#include "janeiro-platform.h"
 #include "mobile-pm.h"
 
 #define SSMT_NS_READ_STREAM_VID_OFFSET(n) (0x1000u + (0x4u * (n)))
@@ -27,16 +27,16 @@ static irqreturn_t janeiro_mailbox_handle_irq(struct edgetpu_dev *etdev,
 					      int irq)
 {
 	struct edgetpu_mailbox *mailbox;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mailbox_manager *mgr = etdev->mailbox_manager;
 	uint i;
-	struct janeiro_platform_dev *jpdev = to_janeiro_dev(etdev);
 
 	if (!mgr)
 		return IRQ_NONE;
-	for (i = 0; i < EDGETPU_NCONTEXTS; i++)
-		if (jpdev->irq[i] == irq)
+	for (i = 0; i < etmdev->n_irq; i++)
+		if (etmdev->irq[i] == irq)
 			break;
-	if (i == EDGETPU_NCONTEXTS)
+	if (i == etmdev->n_irq)
 		return IRQ_NONE;
 	read_lock(&mgr->mailboxes_lock);
 	mailbox = mgr->mailboxes[i];
@@ -70,15 +70,19 @@ u64 edgetpu_chip_tpu_timestamp(struct edgetpu_dev *etdev)
 void edgetpu_chip_init(struct edgetpu_dev *etdev)
 {
 	int i;
-	struct janeiro_platform_dev *jpdev = to_janeiro_dev(etdev);
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 
-	if (!jpdev->ssmt_base)
+	/*
+	 * This only works if the SSMT is set to client-driven mode, which only GSA can do.
+	 * Skip if GSA is not available
+	 */
+	if (!etmdev->ssmt_base || !etmdev->gsa_dev)
 		return;
 
 	/* Setup non-secure SCIDs, assume VID = SCID */
 	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
-		writel(i, SSMT_NS_READ_STREAM_VID_REG(jpdev->ssmt_base, i));
-		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(jpdev->ssmt_base, i));
+		writel(i, SSMT_NS_READ_STREAM_VID_REG(etmdev->ssmt_base, i));
+		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(etmdev->ssmt_base, i));
 	}
 }
 
@@ -120,6 +124,7 @@ int edgetpu_chip_acquire_ext_mailbox(struct edgetpu_client *client,
 		req.count = args->count;
 		req.start = JANEIRO_EXT_DSP_MAILBOX_START;
 		req.end = JANEIRO_EXT_DSP_MAILBOX_END;
+		req.client_type = EDGETPU_EXT_MAILBOX_TYPE_DSP;
 		return edgetpu_mailbox_enable_ext(client, EDGETPU_MAILBOX_ID_USE_ASSOC, &req);
 	}
 	return -ENODEV;
diff --git a/drivers/edgetpu/janeiro-external.c b/drivers/edgetpu/janeiro-external.c
new file mode 100644
index 0000000..c9c3494
--- /dev/null
+++ b/drivers/edgetpu/janeiro-external.c
@@ -0,0 +1,2 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "edgetpu-external.c"
diff --git a/drivers/edgetpu/janeiro-firmware.c b/drivers/edgetpu/janeiro-firmware.c
index ef43931..0339c16 100644
--- a/drivers/edgetpu/janeiro-firmware.c
+++ b/drivers/edgetpu/janeiro-firmware.c
@@ -5,284 +5,4 @@
  * Copyright (C) 2020 Google, Inc.
  */
 
-#include <linux/dma-mapping.h>
-#include <linux/sizes.h>
-#include <linux/slab.h>
-
-#include <linux/iommu.h>
-
-#include "edgetpu.h"
-#include "edgetpu-config.h"
-#include "edgetpu-firmware.h"
-#include "edgetpu-internal.h"
-#include "edgetpu-mmu.h"
-#include "edgetpu-kci.h"
-#include "edgetpu-mailbox.h"
-#include "janeiro-platform.h"
-#include "mobile-firmware.h"
-
-#define MAX_IOMMU_MAPPINGS 26
-
-#define CONFIG_TO_SIZE(a) ((1 << ((a) & 0xFFF)) << 12)
-
-struct iommu_mapping {
-	/* TPU virt address */
-	__u32 virt_address;
-	/*
-	 * contains a 12-bit aligned address and a page-order size into a
-	 * 32-bit value i.e. a physical address and size in page order.
-	 */
-	__u32 image_config_value;
-};
-
-struct janeiro_image_config {
-	__u32 carveout_base;
-	__u32 firmware_base;
-	__u32 firmware_size;
-	struct edgetpu_fw_version firmware_versions;
-	__u32 config_version;
-	__u32 privilege_level;
-	__u32 remapped_region_start;
-	__u32 remapped_region_end;
-	__u32 num_iommu_mapping;
-	struct iommu_mapping mappings[MAX_IOMMU_MAPPINGS];
-} __packed;
-
-struct janeiro_firmware_data {
-	__u32 num_mapping;
-	struct iommu_mapping mappings[MAX_IOMMU_MAPPINGS];
-};
-/*
- * Sets the reset state of the TPU CPU.
- * @val: 1 to put the core in reset state, 0 to release core from reset state.
- */
-static void tpu_cpu_reset(struct edgetpu_dev *etdev, u64 val)
-{
-	edgetpu_dev_write_32_sync(etdev, EDGETPU_REG_RESET_CONTROL, val);
-}
-
-static int janeiro_firmware_after_create(struct edgetpu_firmware *et_fw)
-{
-	struct janeiro_firmware_data *data;
-
-	data = kzalloc(sizeof(*data), GFP_KERNEL);
-	if (!data)
-		return -ENOMEM;
-
-	edgetpu_firmware_set_data(et_fw, data);
-	return 0;
-}
-
-static void janeiro_firmware_before_destroy(struct edgetpu_firmware *et_fw)
-{
-	struct janeiro_firmware_data *data;
-	u32 i, tpu_addr, size;
-	struct edgetpu_dev *etdev = et_fw->etdev;
-
-	if (!etdev->on_exit)
-		tpu_cpu_reset(et_fw->etdev, 1);
-	/* TODO(b/189906347): Remove when GSA/TZ support is available. */
-	/* Remove mappings created by setup_buffer() */
-	data = edgetpu_firmware_get_data(et_fw);
-
-	if (data) {
-		for (i = 0; i < data->num_mapping; i++) {
-			tpu_addr = data->mappings[i].virt_address;
-			size = CONFIG_TO_SIZE(data->mappings[i].image_config_value);
-			edgetpu_mmu_remove_translation(etdev, tpu_addr, size,
-						       EDGETPU_CONTEXT_KCI);
-		}
-		edgetpu_firmware_set_data(et_fw, NULL);
-		kfree(data);
-	}
-}
-
-static int janeiro_firmware_alloc_buffer(struct edgetpu_firmware *et_fw,
-					 struct edgetpu_firmware_buffer *fw_buf)
-{
-	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct janeiro_platform_dev *edgetpu_pdev =
-		container_of(etdev, struct janeiro_platform_dev, edgetpu_dev);
-	/* Allocate extra space for the image header */
-	size_t buffer_size =
-		edgetpu_pdev->fw_region_size + MOBILE_FW_HEADER_SIZE;
-
-	fw_buf->vaddr = kzalloc(buffer_size, GFP_KERNEL);
-	if (!fw_buf->vaddr) {
-		etdev_err(etdev, "%s: failed to allocate buffer (%zu bytes)\n",
-			  __func__, buffer_size);
-		return -ENOMEM;
-	}
-	fw_buf->dma_addr = 0;
-	fw_buf->alloc_size = buffer_size;
-	fw_buf->used_size_align = 16;
-	return 0;
-}
-
-static void janeiro_firmware_free_buffer(struct edgetpu_firmware *et_fw,
-					 struct edgetpu_firmware_buffer *fw_buf)
-{
-	kfree(fw_buf->vaddr);
-	fw_buf->vaddr = NULL;
-	fw_buf->alloc_size = 0;
-	fw_buf->used_size_align = 0;
-}
-
-static int janeiro_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
-					 struct edgetpu_firmware_buffer *fw_buf)
-{
-	int ret = 0;
-	void *image_vaddr;
-	u32 tpu_addr, phys_addr, size, i, j;
-	struct janeiro_image_config *image_config;
-	struct janeiro_firmware_data *data;
-	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct janeiro_platform_dev *edgetpu_pdev =
-		container_of(etdev, struct janeiro_platform_dev, edgetpu_dev);
-
-	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
-		etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
-			  fw_buf->used_size, MOBILE_FW_HEADER_SIZE);
-		return -EINVAL;
-	}
-
-	image_vaddr = memremap(edgetpu_pdev->fw_region_paddr,
-			       edgetpu_pdev->fw_region_size, MEMREMAP_WC);
-	if (!image_vaddr) {
-		etdev_err(etdev, "memremap failed\n");
-		return -ENOMEM;
-	}
-
-	/* fetch the firmware versions */
-	image_config = fw_buf->vaddr + MOBILE_IMAGE_CONFIG_OFFSET;
-	memcpy(&etdev->fw_version, &image_config->firmware_versions,
-	       sizeof(etdev->fw_version));
-
-	/* TODO(b/189906347): Remove when GSA/TZ support is available. */
-	data = edgetpu_firmware_get_data(et_fw);
-	/* Remove old mappings created for previous firmware. */
-	for (i = 0; i < data->num_mapping; i++) {
-		tpu_addr = data->mappings[i].virt_address;
-		size = CONFIG_TO_SIZE(data->mappings[i].image_config_value);
-		phys_addr = (data->mappings[i].image_config_value & ~(0xFFF));
-
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
-	for (i = 0, j = 0; i < MAX_IOMMU_MAPPINGS; i++) {
-		tpu_addr = image_config->mappings[i].virt_address;
-		if (!tpu_addr)
-			continue;
-		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
-		phys_addr = (image_config->mappings[i].image_config_value & ~(0xFFF));
-
-		ret = edgetpu_mmu_add_translation(etdev, tpu_addr, phys_addr, size,
-						  IOMMU_READ | IOMMU_WRITE, EDGETPU_CONTEXT_KCI);
-		if (ret) {
-			etdev_dbg(etdev,
-				  "Unable to map: %d tpu_addr: %#x phys_addr: %pap size: %#x\n",
-				  ret, tpu_addr, &phys_addr, size);
-			goto err;
-		}
-		data->mappings[j].virt_address = tpu_addr;
-		data->mappings[j++].image_config_value =
-						image_config->mappings[i].image_config_value;
-	}
-
-	if (image_config->num_iommu_mapping != j) {
-		etdev_err(etdev, "Invalid firmware header\n");
-		ret = -EINVAL;
-		goto err;
-	}
-	data->num_mapping = j;
-
-	/* Skip the header */
-	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
-	       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
-	memunmap(image_vaddr);
-	return 0;
-err:
-	while (j--) {
-		tpu_addr = data->mappings[j].virt_address;
-		size = CONFIG_TO_SIZE(data->mappings[j].image_config_value);
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
-	data->num_mapping = 0;
-	memunmap(image_vaddr);
-	return ret;
-}
-
-static void
-janeiro_firmware_teardown_buffer(struct edgetpu_firmware *et_fw,
-				 struct edgetpu_firmware_buffer *fw_buf)
-{
-}
-
-static int janeiro_firmware_prepare_run(struct edgetpu_firmware *et_fw,
-					struct edgetpu_firmware_buffer *fw_buf)
-{
-	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct janeiro_platform_dev *edgetpu_pdev =
-		container_of(etdev, struct janeiro_platform_dev, edgetpu_dev);
-	dma_addr_t fw_dma_addr;
-	int ret = 0;
-
-	fw_dma_addr = edgetpu_pdev->fw_region_paddr;
-
-	//TODO: enable as and when sysmmu started working correctly on hybrid
-	//platform
-#if 0
-	/* Clear Substream ID (aka SCID) for instruction remapped addresses */
-	u32 sec_reg = edgetpu_dev_read_32(
-			etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY);
-	sec_reg &= ~(0x0F << 16);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY,
-			sec_reg);
-
-	/* Clear Substream ID (aka SCID) for all other addresses */
-	sec_reg = edgetpu_dev_read_32(etdev, EDGETPU_REG_SECURITY);
-	sec_reg &= ~(0x0F << 16);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_SECURITY, sec_reg);
-//	dma_sync_single_for_device(etdev->dev, fw_dma_addr,
-//				   fw_buf->used_size - JANEIRO_FW_HEADER_SIZE,
-//				   DMA_TO_DEVICE);
-#endif
-	tpu_cpu_reset(etdev, 1);
-
-	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
-	edgetpu_mailbox_reset(etdev->kci->mailbox);
-
-	tpu_cpu_reset(etdev, 0);
-	//TODO: cleanup
-	return ret;
-}
-
-static const struct edgetpu_firmware_chip_data janeiro_firmware_chip_data = {
-	.default_firmware_name = EDGETPU_DEFAULT_FIRMWARE_NAME,
-	.after_create = janeiro_firmware_after_create,
-	.before_destroy = janeiro_firmware_before_destroy,
-	.alloc_buffer = janeiro_firmware_alloc_buffer,
-	.free_buffer = janeiro_firmware_free_buffer,
-	.setup_buffer = janeiro_firmware_setup_buffer,
-	.teardown_buffer = janeiro_firmware_teardown_buffer,
-	.prepare_run = janeiro_firmware_prepare_run,
-};
-
-int mobile_edgetpu_firmware_create(struct edgetpu_dev *etdev)
-{
-	return edgetpu_firmware_create(etdev, &janeiro_firmware_chip_data);
-}
-
-void mobile_edgetpu_firmware_destroy(struct edgetpu_dev *etdev)
-{
-	edgetpu_firmware_destroy(etdev);
-}
-
-unsigned long edgetpu_chip_firmware_iova(struct edgetpu_dev *etdev)
-{
-	/*
-	 * There is no IOVA in Janeiro, since firmware the IOMMU is
-	 * bypassed and the only translation in effect is the one
-	 * done by instruction remap registers
-	 */
-	return EDGETPU_INSTRUCTION_REMAP_BASE;
-}
+#include "mobile-firmware.c"
diff --git a/drivers/edgetpu/janeiro-platform.c b/drivers/edgetpu/janeiro-platform.c
index 5708a5d..9ca3ff3 100644
--- a/drivers/edgetpu/janeiro-platform.c
+++ b/drivers/edgetpu/janeiro-platform.c
@@ -6,22 +6,18 @@
  */
 
 #include <linux/device.h>
+#include <linux/io.h>
+#include <linux/mod_devicetable.h>
 #include <linux/module.h>
-#include <linux/iopoll.h>
 #include <linux/of.h>
-#include <linux/of_address.h>
 #include <linux/platform_device.h>
 
 #include "edgetpu-config.h"
-#include "edgetpu-debug-dump.h"
-#include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
-#include "edgetpu-iremap-pool.h"
-#include "edgetpu-mmu.h"
-#include "edgetpu-telemetry.h"
+#include "edgetpu-mobile-platform.h"
 #include "janeiro-platform.h"
-#include "mobile-firmware.h"
-#include "mobile-pm.h"
+
+#include "edgetpu-mobile-platform.c"
 
 static const struct of_device_id edgetpu_of_match[] = {
 	/* TODO(b/190677977): remove  */
@@ -32,117 +28,6 @@ static const struct of_device_id edgetpu_of_match[] = {
 
 MODULE_DEVICE_TABLE(of, edgetpu_of_match);
 
-#define EDGETPU_POOL_MEM_OFFSET (EDGETPU_TELEMETRY_BUFFER_SIZE * 2)
-
-static void janeiro_get_telemetry_mem(struct janeiro_platform_dev *etpdev,
-				      enum edgetpu_telemetry_type type,
-				      struct edgetpu_coherent_mem *mem)
-{
-	int offset = type == EDGETPU_TELEMETRY_TRACE ?
-				   EDGETPU_TELEMETRY_BUFFER_SIZE :
-				   0;
-	mem->vaddr = etpdev->shared_mem_vaddr + offset;
-	mem->dma_addr = EDGETPU_REMAPPED_DATA_ADDR + offset;
-	mem->tpu_addr = EDGETPU_REMAPPED_DATA_ADDR + offset;
-	mem->host_addr = 0;
-	mem->size = EDGETPU_TELEMETRY_BUFFER_SIZE;
-}
-
-static int janeiro_platform_setup_fw_region(struct janeiro_platform_dev *etpdev)
-{
-	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
-	struct device *dev = etdev->dev;
-	struct resource r;
-	struct device_node *np;
-	int err;
-	size_t region_map_size =
-		EDGETPU_FW_SIZE_MAX + EDGETPU_REMAPPED_DATA_SIZE;
-
-	np = of_parse_phandle(dev->of_node, "memory-region", 0);
-	if (!np) {
-		dev_err(dev, "No memory region for firmware\n");
-		return -ENODEV;
-	}
-
-	err = of_address_to_resource(np, 0, &r);
-	of_node_put(np);
-	if (err) {
-		dev_err(dev, "No memory address assigned to firmware region\n");
-		return err;
-	}
-
-	if (resource_size(&r) < region_map_size) {
-		dev_err(dev,
-			"Memory region for firmware too small (%zu bytes needed, got %llu)\n",
-			region_map_size, resource_size(&r));
-		return -ENOSPC;
-	}
-
-	etpdev->fw_region_paddr = r.start;
-	etpdev->fw_region_size = EDGETPU_FW_SIZE_MAX;
-
-	etpdev->shared_mem_vaddr =
-		memremap(r.start + EDGETPU_REMAPPED_DATA_OFFSET,
-			 EDGETPU_REMAPPED_DATA_SIZE, MEMREMAP_WC);
-	if (!etpdev->shared_mem_vaddr) {
-		dev_err(dev, "Shared memory remap failed\n");
-		return -EINVAL;
-	}
-	etpdev->shared_mem_paddr = r.start + EDGETPU_REMAPPED_DATA_OFFSET;
-
-	return 0;
-}
-
-static void
-janeiro_platform_cleanup_fw_region(struct janeiro_platform_dev *etpdev)
-{
-	if (!etpdev->shared_mem_vaddr)
-		return;
-	memunmap(etpdev->shared_mem_vaddr);
-	etpdev->shared_mem_vaddr = NULL;
-}
-
-int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev)
-{
-	int ret;
-
-	ret = edgetpu_mmu_attach(etdev, NULL);
-	if (ret)
-		dev_err(etdev->dev, "failed to attach IOMMU: %d\n", ret);
-	return ret;
-}
-
-void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
-{
-	edgetpu_mmu_detach(etdev);
-}
-
-static int janeiro_parse_ssmt(struct janeiro_platform_dev *etpdev)
-{
-	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
-	struct platform_device *pdev = to_platform_device(etdev->dev);
-	struct resource *res;
-	int rc;
-	void __iomem *ssmt_base;
-
-	/* TODO(b/197301774): Remove when GSA configure SSMT */
-	return -EINVAL;
-
-	res = platform_get_resource_byname(pdev, IORESOURCE_MEM, "ssmt");
-	if (!res) {
-		etdev_warn(etdev, "Failed to find SSMT register base");
-		return -EINVAL;
-	}
-	ssmt_base = devm_ioremap_resource(&pdev->dev, res);
-	if (IS_ERR(ssmt_base)) {
-		rc = PTR_ERR(ssmt_base);
-		etdev_warn(etdev, "Failed to map SSMT register base: %d\n", rc);
-		return rc;
-	}
-	etpdev->ssmt_base = ssmt_base;
-	return 0;
-}
-
 /*
  * Set shareability for enabling IO coherency in Janeiro
  */
@@ -162,207 +47,49 @@ static int janeiro_mmu_set_shareability(struct device *dev, u32 reg_base)
 	return 0;
 }
 
-static int janeiro_parse_dt(struct device *dev)
+static int janeiro_parse_set_dt_property(struct device *dev)
 {
 	int ret;
 	u32 reg;
 
-	if (of_find_property(dev->of_node, "edgetpu,shareability", NULL)) {
-		ret = of_property_read_u32_index(dev->of_node, "edgetpu,shareability", 0, &reg);
-		if (ret)
-			return ret;
-	} else {
-		/*
-		 * TODO(b/193593081): Remove compatibility code
-		 * Fallback for older driver versions until DT property support is
-		 * widely adopted.
-		 */
-		reg = EDGETPU_SYSREG_TPU_BASE;
-	}
-
+	if (!of_find_property(dev->of_node, "edgetpu,shareability", NULL))
+		return -ENODEV;
+	ret = of_property_read_u32_index(dev->of_node, "edgetpu,shareability", 0, &reg);
+	if (ret)
+		return ret;
 	return janeiro_mmu_set_shareability(dev, reg);
 }
 
-static int edgetpu_platform_probe(struct platform_device *pdev)
+static int janeiro_platform_after_probe(struct edgetpu_mobile_platform_dev *etmdev)
 {
-	struct device *dev = &pdev->dev;
-	struct janeiro_platform_dev *edgetpu_pdev;
-	struct resource *r;
-	struct edgetpu_mapped_resource regs;
-	int ret, i;
-	struct edgetpu_iface_params iface_params[] = {
-		/* Default interface */
-		{ .name = NULL },
-		/* Common name for SoC embedded devices */
-		{ .name = "edgetpu-soc" },
-	};
-
-	edgetpu_pdev = devm_kzalloc(dev, sizeof(*edgetpu_pdev), GFP_KERNEL);
-	if (!edgetpu_pdev)
-		return -ENOMEM;
-
-	platform_set_drvdata(pdev, &edgetpu_pdev->edgetpu_dev);
-	edgetpu_pdev->edgetpu_dev.dev = dev;
-
-	r = platform_get_resource(pdev, IORESOURCE_MEM, 0);
-	if (IS_ERR_OR_NULL(r)) {
-		dev_err(dev, "failed to get memory resource\n");
-		return -ENODEV;
-	}
-
-	regs.phys = r->start;
-	regs.size = resource_size(r);
-
-	regs.mem = devm_ioremap_resource(dev, r);
-	if (IS_ERR_OR_NULL(regs.mem)) {
-		dev_err(dev, "failed to map registers\n");
-		return -ENODEV;
-	}
-
-	ret = mobile_pm_create(&edgetpu_pdev->edgetpu_dev);
-	if (ret) {
-		dev_err(dev, "Failed to initialize PM interface (%d)\n", ret);
-		return ret;
-	}
-
-	ret = janeiro_platform_setup_fw_region(edgetpu_pdev);
-	if (ret) {
-		dev_err(dev, "%s setup fw regions failed: %d\n", DRIVER_NAME,
-			ret);
-		goto out;
-	}
-
-	ret = edgetpu_iremap_pool_create(
-		&edgetpu_pdev->edgetpu_dev,
-		/* Base virtual address (kernel address space) */
-		edgetpu_pdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
-		/* Base DMA address */
-		EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
-		/* Base TPU address */
-		EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
-		/* Base physical address */
-		edgetpu_pdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
-		/* Size */
-		EDGETPU_REMAPPED_DATA_SIZE - EDGETPU_POOL_MEM_OFFSET,
-		/* Granularity */
-		PAGE_SIZE);
-	if (ret) {
-		dev_err(dev,
-			"%s failed to initialize remapped memory pool: %d\n",
-			DRIVER_NAME, ret);
-		goto out_cleanup_fw;
-	}
-
-	edgetpu_pdev->edgetpu_dev.mcp_id = -1;
-	edgetpu_pdev->edgetpu_dev.mcp_die_index = 0;
-
-	for (i = 0; i < EDGETPU_NCONTEXTS; i++)
-		edgetpu_pdev->irq[i] = platform_get_irq(pdev, i);
-
-	ret = janeiro_parse_dt(dev);
-
-	if (ret)
-		dev_warn(dev, "%s failed to enable shareability: %d\n",
-			 DRIVER_NAME, ret);
-
-	ret = edgetpu_device_add(&edgetpu_pdev->edgetpu_dev, &regs, iface_params,
-				 ARRAY_SIZE(iface_params));
-
-	if (ret) {
-		dev_err(dev, "%s edgetpu setup failed: %d\n", DRIVER_NAME, ret);
-		goto out_destroy_iremap;
-	}
-	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
-		if (edgetpu_pdev->irq[i] >= 0)
-			ret = edgetpu_register_irq(&edgetpu_pdev->edgetpu_dev,
-						   edgetpu_pdev->irq[i]);
-		if (ret)
-			break;
-	}
-	if (ret) {
-		while (i-- > 0) {
-			edgetpu_unregister_irq(&edgetpu_pdev->edgetpu_dev,
-					       edgetpu_pdev->irq[i]);
-		}
-		dev_err(dev, "%s edgetpu irq registration failed: %d\n",
-			DRIVER_NAME, ret);
-		goto out_remove_device;
-	}
-
-	ret = janeiro_parse_ssmt(edgetpu_pdev);
-	if (ret)
-		dev_warn(
-			dev,
-			"SSMT setup failed (%d). Context isolation not enforced\n",
-			ret);
-
-	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_LOG,
-				  &edgetpu_pdev->log_mem);
-	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_TRACE,
-				  &edgetpu_pdev->trace_mem);
+	int ret;
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
 
-	ret = edgetpu_telemetry_init(&edgetpu_pdev->edgetpu_dev,
-				     &edgetpu_pdev->log_mem,
-				     &edgetpu_pdev->trace_mem);
+	ret = janeiro_parse_set_dt_property(etdev->dev);
 	if (ret)
-		goto out_remove_device;
+		dev_warn(etdev->dev, "failed to enable shareability: %d", ret);
 
-	ret = mobile_edgetpu_firmware_create(&edgetpu_pdev->edgetpu_dev);
-	if (ret) {
-		dev_err(dev, "%s initialize firmware downloader failed: %d\n",
-			DRIVER_NAME, ret);
-		goto out_tel_exit;
-	}
+	return 0;
+}
 
-	dev_info(dev, "%s edgetpu initialized. Build: %s\n",
-		 edgetpu_pdev->edgetpu_dev.dev_name, GIT_REPO_TAG);
+static int edgetpu_platform_probe(struct platform_device *pdev)
+{
+	struct device *dev = &pdev->dev;
+	struct janeiro_platform_dev *jpdev;
+	struct edgetpu_mobile_platform_dev *etmdev;
 
-	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, false);
-out:
-	dev_dbg(dev, "Probe finished\n");
+	jpdev = devm_kzalloc(dev, sizeof(*jpdev), GFP_KERNEL);
+	if (!jpdev)
+		return -ENOMEM;
 
-	return 0;
-out_tel_exit:
-	edgetpu_telemetry_exit(&edgetpu_pdev->edgetpu_dev);
-out_remove_device:
-	edgetpu_device_remove(&edgetpu_pdev->edgetpu_dev);
-out_destroy_iremap:
-	edgetpu_iremap_pool_destroy(&edgetpu_pdev->edgetpu_dev);
-out_cleanup_fw:
-	janeiro_platform_cleanup_fw_region(edgetpu_pdev);
-	dev_dbg(dev, "Probe finished with error %d, powering down\n", ret);
-	return ret;
+	etmdev = &jpdev->mobile_dev;
+	etmdev->after_probe = janeiro_platform_after_probe;
+	return edgetpu_mobile_platform_probe(pdev, etmdev);
 }
 
 static int edgetpu_platform_remove(struct platform_device *pdev)
 {
-	int i;
-	struct edgetpu_dev *etdev = platform_get_drvdata(pdev);
-	struct janeiro_platform_dev *janeiro_pdev = to_janeiro_dev(etdev);
-
-	/* TODO(b/189906347): Use edgetpu_device_remove() for cleanup after
-	 * having GSA/TZ support.
-	 */
-	etdev->on_exit = true;
-	edgetpu_pm_get(etdev->pm);
-	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
-		if (janeiro_pdev->irq[i] >= 0)
-			edgetpu_unregister_irq(etdev, janeiro_pdev->irq[i]);
-	}
-	mobile_edgetpu_firmware_destroy(etdev);
-	edgetpu_telemetry_exit(etdev);
-	edgetpu_chip_exit(etdev);
-	edgetpu_debug_dump_exit(etdev);
-	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
-	edgetpu_pm_put(etdev->pm);
-	edgetpu_pm_shutdown(etdev, true);
-	edgetpu_usage_stats_exit(etdev);
-	edgetpu_chip_remove_mmu(etdev);
-	edgetpu_fs_remove(etdev);
-	edgetpu_iremap_pool_destroy(etdev);
-	janeiro_platform_cleanup_fw_region(janeiro_pdev);
-	mobile_pm_destroy(etdev);
-	return 0;
+	return edgetpu_mobile_platform_remove(pdev);
 }
 
 static struct platform_driver edgetpu_platform_driver = {
@@ -390,7 +117,7 @@ static void __exit edgetpu_platform_exit(void)
 	edgetpu_exit();
 }
 
-MODULE_DESCRIPTION("Janeiro Edge TPU platform driver");
+MODULE_DESCRIPTION("Google Edge TPU platform driver");
 MODULE_LICENSE("GPL v2");
 module_init(edgetpu_platform_init);
 module_exit(edgetpu_platform_exit);
diff --git a/drivers/edgetpu/janeiro-platform.h b/drivers/edgetpu/janeiro-platform.h
index 51eb8d6..5ffe734 100644
--- a/drivers/edgetpu/janeiro-platform.h
+++ b/drivers/edgetpu/janeiro-platform.h
@@ -4,50 +4,18 @@
  *
  * Copyright (C) 2020 Google, Inc.
  */
+
 #ifndef __JANEIRO_PLATFORM_H__
 #define __JANEIRO_PLATFORM_H__
 
-#include <linux/device.h>
-#include <linux/io.h>
-#include <linux/mutex.h>
-#include <linux/types.h>
-
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-#include <soc/google/bcl.h>
-#endif
-
 #include "edgetpu-internal.h"
+#include "edgetpu-mobile-platform.h"
 
-#define to_janeiro_dev(etdev)                                                  \
-	container_of(etdev, struct janeiro_platform_dev, edgetpu_dev)
-
-// TODO(b/176881607): merge with abrolhos
-struct janeiro_platform_pwr {
-	struct mutex state_lock;
-	u64 min_state;
-	u64 requested_state;
-};
+#define to_janeiro_dev(etdev)                                                                      \
+	container_of((to_mobile_dev(etdev)), struct janeiro_platform_dev, mobile_dev)
 
 struct janeiro_platform_dev {
-	struct edgetpu_dev edgetpu_dev;
-	struct janeiro_platform_pwr platform_pwr;
-	int irq[EDGETPU_NCONTEXTS];
-	phys_addr_t fw_region_paddr;
-	void *fw_region_vaddr;
-	size_t fw_region_size;
-	void *shared_mem_vaddr;
-	phys_addr_t shared_mem_paddr;
-	size_t shared_mem_size;
-	phys_addr_t csr_paddr;
-	dma_addr_t csr_iova;
-	size_t csr_size;
-	struct device *gsa_dev;
-	void __iomem *ssmt_base;
-	struct edgetpu_coherent_mem log_mem;
-	struct edgetpu_coherent_mem trace_mem;
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-	struct bcl_device *bcl_dev;
-#endif
+	struct edgetpu_mobile_platform_dev mobile_dev;
 };
 
 #endif /* __JANEIRO_PLATFORM_H__ */
diff --git a/drivers/edgetpu/janeiro-pm.c b/drivers/edgetpu/janeiro-pm.c
index 700601e..681f58d 100644
--- a/drivers/edgetpu/janeiro-pm.c
+++ b/drivers/edgetpu/janeiro-pm.c
@@ -5,26 +5,17 @@
  * Copyright (C) 2020 Google, Inc.
  */
 
-#include <linux/atomic.h>
 #include <linux/delay.h>
 #include <linux/iopoll.h>
-#include <linux/module.h>
-#include <linux/pm_runtime.h>
-#include <linux/version.h>
 
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-#include <soc/google/bcl.h>
-#endif
-
-#include "edgetpu-firmware.h"
+#include "edgetpu-config.h"
 #include "edgetpu-internal.h"
-#include "edgetpu-kci.h"
-#include "edgetpu-mailbox.h"
-#include "edgetpu-pm.h"
-#include "janeiro-platform.h"
+#include "edgetpu-mobile-platform.h"
 #include "mobile-pm.h"
 
-#include "edgetpu-pm.c"
+#define TPU_DEFAULT_POWER_STATE		TPU_ACTIVE_NOM
+
+#include "mobile-pm.c"
 
 #define SHUTDOWN_DELAY_US_MIN		20
 #define SHUTDOWN_DELAY_US_MAX		20
@@ -32,200 +23,6 @@
 #define BOOTUP_DELAY_US_MAX		250
 #define SHUTDOWN_MAX_DELAY_COUNT	50
 
-/* Default power state */
-static int power_state = TPU_ACTIVE_NOM;
-
-module_param(power_state, int, 0660);
-
-static struct dentry *janeiro_pwr_debugfs_dir;
-
-enum edgetpu_pwr_state edgetpu_active_states[EDGETPU_NUM_STATES] = {
-	TPU_ACTIVE_UUD,
-	TPU_ACTIVE_SUD,
-	TPU_ACTIVE_UD,
-	TPU_ACTIVE_NOM,
-};
-
-uint32_t *edgetpu_states_display = edgetpu_active_states;
-
-static int janeiro_pwr_state_init(struct device *dev)
-{
-	int ret;
-	int curr_state;
-
-	pm_runtime_enable(dev);
-	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
-
-	if (curr_state > TPU_OFF) {
-		ret = pm_runtime_get_sync(dev);
-		if (ret) {
-			dev_err(dev, "%s: pm_runtime_get_sync err: %d\n",
-				__func__, ret);
-			return ret;
-		}
-	}
-	ret = exynos_acpm_set_init_freq(TPU_ACPM_DOMAIN, curr_state);
-	if (ret) {
-		dev_err(dev, "error initializing tpu ACPM freq: %d\n", ret);
-		if (curr_state > TPU_OFF)
-			pm_runtime_put_sync(dev);
-		return ret;
-	}
-	return ret;
-}
-
-static int janeiro_pwr_state_set_locked(void *data, u64 val)
-{
-	int ret;
-	int curr_state;
-	int timeout_cnt = 0;
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct device *dev = etdev->dev;
-
-	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
-
-	dev_dbg(dev, "Power state %d -> %llu\n", curr_state, val);
-
-	if (curr_state == TPU_OFF && val > TPU_OFF) {
-		ret = pm_runtime_get_sync(dev);
-		if (ret) {
-			dev_err(dev, "%s: pm_runtime_get_sync err: %d\n",
-				__func__, ret);
-			return ret;
-		}
-	}
-
-	/* TPU_OFF is invalid state */
-	if (val != TPU_OFF) {
-		ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, (unsigned long)val);
-		if (ret) {
-			dev_err(dev, "error setting tpu power state: %d\n", ret);
-			pm_runtime_put_sync(dev);
-			return ret;
-		}
-	}
-
-	if (curr_state != TPU_OFF && val == TPU_OFF) {
-		ret = pm_runtime_put_sync(dev);
-		if (ret) {
-			dev_err(dev, "%s: pm_runtime_put_sync returned %d\n",
-				__func__, ret);
-			return ret;
-		}
-		do {
-			/* Delay 20us per retry till blk shutdown finished */
-			usleep_range(SHUTDOWN_DELAY_US_MIN, SHUTDOWN_DELAY_US_MAX);
-			/* Only poll for BLK status instead of CLK rate */
-			curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 1);
-			if (!curr_state)
-				break;
-			timeout_cnt++;
-		} while (timeout_cnt < SHUTDOWN_MAX_DELAY_COUNT);
-		if (timeout_cnt == SHUTDOWN_MAX_DELAY_COUNT)
-			dev_warn(dev, "%s: blk_shutdown timeout\n", __func__);
-	}
-
-	return ret;
-}
-
-static int janeiro_pwr_state_get_locked(void *data, u64 *val)
-{
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct device *dev = etdev->dev;
-
-	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
-	dev_dbg(dev, "current tpu power state: %llu\n", *val);
-
-	return 0;
-}
-
-static int janeiro_pwr_state_set(void *data, u64 val)
-{
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
-	int ret = 0;
-
-	mutex_lock(&platform_pwr->state_lock);
-	platform_pwr->requested_state = val;
-	if (val >= platform_pwr->min_state)
-		ret = janeiro_pwr_state_set_locked(etdev, val);
-	mutex_unlock(&platform_pwr->state_lock);
-	return ret;
-}
-
-static int janeiro_pwr_state_get(void *data, u64 *val)
-{
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
-	int ret;
-
-	mutex_lock(&platform_pwr->state_lock);
-	ret = janeiro_pwr_state_get_locked(etdev, val);
-	mutex_unlock(&platform_pwr->state_lock);
-	return ret;
-}
-
-static int janeiro_min_pwr_state_set(void *data, u64 val)
-{
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
-	int ret = 0;
-
-	mutex_lock(&platform_pwr->state_lock);
-	platform_pwr->min_state = val;
-	if (val >= platform_pwr->requested_state)
-		ret = janeiro_pwr_state_set_locked(etdev, val);
-	mutex_unlock(&platform_pwr->state_lock);
-	return ret;
-}
-
-static int janeiro_min_pwr_state_get(void *data, u64 *val)
-{
-	struct edgetpu_dev *etdev = (typeof(etdev))data;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
-
-	mutex_lock(&platform_pwr->state_lock);
-	*val = platform_pwr->min_state;
-	mutex_unlock(&platform_pwr->state_lock);
-	return 0;
-}
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, janeiro_pwr_state_get,
-			 janeiro_pwr_state_set, "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, janeiro_min_pwr_state_get,
-			 janeiro_min_pwr_state_set, "%llu\n");
-
-static int janeiro_get_initial_pwr_state(struct device *dev)
-{
-	switch (power_state) {
-	case TPU_ACTIVE_UUD:
-	case TPU_ACTIVE_SUD:
-	case TPU_ACTIVE_UD:
-	case TPU_ACTIVE_NOM:
-		dev_info(dev, "Initial power state: %d\n", power_state);
-		break;
-	case TPU_OFF:
-		dev_warn(dev, "Power state %d prevents control core booting",
-			 power_state);
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 4, 0)
-		fallthrough;
-#endif
-	default:
-		dev_warn(dev, "Power state %d is invalid\n", power_state);
-		dev_warn(dev, "defaulting to active nominal\n");
-		power_state = TPU_ACTIVE_NOM;
-		break;
-	}
-	return power_state;
-}
-
-static void janeiro_power_down(struct edgetpu_pm *etpm);
-
 #define EDGETPU_PSM0_CFG 0x1c1880
 #define EDGETPU_PSM0_START 0x1c1884
 #define EDGETPU_PSM0_STATUS 0x1c1888
@@ -234,21 +31,39 @@ static void janeiro_power_down(struct edgetpu_pm *etpm);
 #define EDGETPU_PSM1_STATUS 0x1c2888
 #define EDGETPU_LPM_CHANGE_TIMEOUT 30000
 
-static int janeiro_set_lpm(struct edgetpu_dev *etdev)
+static void janeiro_lpm_down(struct edgetpu_dev *etdev)
+{
+	int timeout_cnt = 0;
+	u32 val;
+
+	do {
+		/* Manually delay 20us per retry till LPM shutdown finished */
+		usleep_range(SHUTDOWN_DELAY_US_MIN, SHUTDOWN_DELAY_US_MAX);
+		val = edgetpu_dev_read_32_sync(etdev, EDGETPU_REG_LPM_CONTROL);
+		if ((val & 0x1000) || (val == 0))
+			break;
+		timeout_cnt++;
+	} while (timeout_cnt < SHUTDOWN_MAX_DELAY_COUNT);
+	if (timeout_cnt == SHUTDOWN_MAX_DELAY_COUNT)
+		// Log the issue then continue to perform the shutdown forcefully.
+		etdev_warn(etdev, "LPM shutdown failure, continuing BLK shutdown\n");
+}
+
+static int janeiro_lpm_up(struct edgetpu_dev *etdev)
 {
 	int ret;
 	u32 val;
 
 	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_START, 1);
-	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM0_STATUS, val,
-				 val & 0x80, 20, EDGETPU_LPM_CHANGE_TIMEOUT);
+	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM0_STATUS, val, val & 0x80, 20,
+				 EDGETPU_LPM_CHANGE_TIMEOUT);
 	if (ret) {
 		etdev_err(etdev, "Set LPM0 failed: %d\n", ret);
 		return ret;
 	}
 	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM1_START, 1);
-	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM1_STATUS, val,
-				 val & 0x80, 20, EDGETPU_LPM_CHANGE_TIMEOUT);
+	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM1_STATUS, val, val & 0x80, 20,
+				 EDGETPU_LPM_CHANGE_TIMEOUT);
 	if (ret) {
 		etdev_err(etdev, "Set LPM1 failed: %d\n", ret);
 		return ret;
@@ -260,194 +75,50 @@ static int janeiro_set_lpm(struct edgetpu_dev *etdev)
 	return 0;
 }
 
-static int janeiro_power_up(struct edgetpu_pm *etpm)
-{
-	struct edgetpu_dev *etdev = etpm->etdev;
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-#endif
-	int ret = 0;
-
-	ret = janeiro_pwr_state_set(
-		etpm->etdev, janeiro_get_initial_pwr_state(etdev->dev));
-
-	etdev_info(etpm->etdev, "Powering up\n");
-
-	if (ret)
-		return ret;
-
-	/* Delay 100us to make sure LPM is accessible */
-	usleep_range(BOOTUP_DELAY_US_MIN, BOOTUP_DELAY_US_MAX);
-	janeiro_set_lpm(etdev);
-
-	edgetpu_chip_init(etdev);
-
-	if (etdev->kci) {
-		etdev_dbg(etdev, "Resetting KCI\n");
-		edgetpu_kci_reinit(etdev->kci);
-	}
-	if (etdev->mailbox_manager) {
-		etdev_dbg(etdev, "Resetting VII mailboxes\n");
-		edgetpu_mailbox_reset_vii(etdev->mailbox_manager);
-	}
-
-	if (!etdev->firmware || etdev->on_exit)
-		return 0;
-
-	/*
-	 * Why this function uses edgetpu_firmware_*_locked functions without explicitly holding
-	 * edgetpu_firmware_lock:
-	 *
-	 * edgetpu_pm_get() is called in two scenarios - one is when the firmware loading is
-	 * attempting, another one is when the user-space clients need the device be powered
-	 * (usually through acquiring the wakelock).
-	 *
-	 * For the first scenario edgetpu_firmware_is_loading() below shall return true.
-	 * For the second scenario we are indeed called without holding the firmware lock, but the
-	 * firmware loading procedures (i.e. the first scenario) always call edgetpu_pm_get() before
-	 * changing the firmware state, and edgetpu_pm_get() is blocked until this function
-	 * finishes. In short, we are protected by the PM lock.
-	 */
-
-	if (edgetpu_firmware_is_loading(etdev))
-		return 0;
-
-	/* attempt firmware run */
-	switch (edgetpu_firmware_status_locked(etdev)) {
-	case FW_VALID:
-		ret = edgetpu_firmware_restart_locked(etdev, false);
-		break;
-	case FW_INVALID:
-		ret = edgetpu_firmware_run_default_locked(etdev);
-		break;
-	default:
-		break;
-	}
-
-	if (ret) {
-		janeiro_power_down(etpm);
-	} else {
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-		if (!edgetpu_pdev->bcl_dev)
-			edgetpu_pdev->bcl_dev = google_retrieve_bcl_handle();
-		if (edgetpu_pdev->bcl_dev)
-			google_init_tpu_ratio(edgetpu_pdev->bcl_dev);
-#endif
-	}
-
-	return ret;
-}
-
-static void
-janeiro_pm_shutdown_firmware(struct janeiro_platform_dev *etpdev,
-			     struct edgetpu_dev *etdev)
-{
-	int ret;
-
-	ret = edgetpu_kci_shutdown(etdev->kci);
-	if (ret) {
-		etdev_err(etdev, "firmware shutdown failed: %d",
-			  ret);
-		return;
-	}
-}
-
-static void janeiro_power_down(struct edgetpu_pm *etpm)
+static void janeiro_block_down(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_dev *etdev = etpm->etdev;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	u64 val;
 	int timeout_cnt = 0;
+	int curr_state;
 
-	etdev_info(etdev, "Powering down\n");
-
-	if (janeiro_pwr_state_get(etdev, &val)) {
-		etdev_warn(etdev, "Failed to read current power state\n");
-		val = TPU_ACTIVE_NOM;
-	}
-	if (val == TPU_OFF) {
-		etdev_dbg(etdev, "Device already off, skipping shutdown\n");
-		return;
-	}
-
-	if (etdev->kci && edgetpu_firmware_status_locked(etdev) == FW_VALID) {
-		/* Update usage stats before we power off fw. */
-		edgetpu_kci_update_usage_locked(etdev);
-		janeiro_pm_shutdown_firmware(edgetpu_pdev, etdev);
-		edgetpu_kci_cancel_work_queues(etdev->kci);
-	}
 	do {
-		/* Manually delay 20us per retry till LPM shutdown finished */
+		/* Delay 20us per retry till blk shutdown finished */
 		usleep_range(SHUTDOWN_DELAY_US_MIN, SHUTDOWN_DELAY_US_MAX);
-		val = edgetpu_dev_read_32_sync(etdev, EDGETPU_REG_LPM_CONTROL);
-		if ((val & 0x1000) || (val == 0))
+		/* Only poll for BLK status instead of CLK rate */
+		curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 1);
+		if (!curr_state)
 			break;
 		timeout_cnt++;
 	} while (timeout_cnt < SHUTDOWN_MAX_DELAY_COUNT);
 	if (timeout_cnt == SHUTDOWN_MAX_DELAY_COUNT)
-		// Log the issue then continue to perform the shutdown forcefully.
-		etdev_err(etdev, "LPM shutdown failure, performing BLK shutdown\n");
-	janeiro_pwr_state_set(etdev, TPU_OFF);
+		etdev_warn(etdev, "%s: blk_shutdown timeout\n", __func__);
 }
 
-static int janeiro_pm_after_create(struct edgetpu_pm *etpm)
+static void janeiro_firmware_down(struct edgetpu_dev *etdev)
 {
 	int ret;
-	struct edgetpu_dev *etdev = etpm->etdev;
-	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
-	struct device *dev = etdev->dev;
-
-	ret = janeiro_pwr_state_init(dev);
-	if (ret)
-		return ret;
 
-	mutex_init(&edgetpu_pdev->platform_pwr.state_lock);
-
-	ret = janeiro_pwr_state_set(etdev,
-				     janeiro_get_initial_pwr_state(dev));
-	if (ret)
-		return ret;
-	janeiro_pwr_debugfs_dir =
-		debugfs_create_dir("power", edgetpu_fs_debugfs_dir());
-	if (IS_ERR_OR_NULL(janeiro_pwr_debugfs_dir)) {
-		etdev_warn(etdev, "Failed to create debug FS power");
-		/* don't fail the procedure on debug FS creation fails */
-		return 0;
+	ret = edgetpu_kci_shutdown(etdev->kci);
+	if (ret) {
+		etdev_err(etdev, "firmware shutdown failed: %d", ret);
+		return;
 	}
-	debugfs_create_file("state", 0660, janeiro_pwr_debugfs_dir, etdev,
-			    &fops_tpu_pwr_state);
-	debugfs_create_file("min_state", 0660, janeiro_pwr_debugfs_dir, etdev,
-			    &fops_tpu_min_pwr_state);
-	return 0;
 }
 
-static void janeiro_pm_before_destroy(struct edgetpu_pm *etpm)
+static int janeiro_acpm_set_rate(unsigned int id, unsigned long rate)
 {
-	debugfs_remove_recursive(janeiro_pwr_debugfs_dir);
-	pm_runtime_disable(etpm->etdev->dev);
+	return exynos_acpm_set_rate(id, rate);
 }
 
-static struct edgetpu_pm_handlers janeiro_pm_handlers = {
-	.after_create = janeiro_pm_after_create,
-	.before_destroy = janeiro_pm_before_destroy,
-	.power_up = janeiro_power_up,
-	.power_down = janeiro_power_down,
-};
-
-int mobile_pm_create(struct edgetpu_dev *etdev)
-{
-	return edgetpu_pm_create(etdev, &janeiro_pm_handlers);
-}
-
-void mobile_pm_destroy(struct edgetpu_dev *etdev)
+int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 {
-	edgetpu_pm_destroy(etdev);
-}
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
-{
-}
+	platform_pwr->lpm_up = janeiro_lpm_up;
+	platform_pwr->lpm_down = janeiro_lpm_down;
+	platform_pwr->block_down = janeiro_block_down;
+	platform_pwr->firmware_down = janeiro_firmware_down;
+	platform_pwr->acpm_set_rate = janeiro_acpm_set_rate;
 
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val)
-{
+	return mobile_pm_create(etdev);
 }
diff --git a/drivers/edgetpu/janeiro/config-pwr-state.h b/drivers/edgetpu/janeiro/config-pwr-state.h
index cdea172..f30d82e 100644
--- a/drivers/edgetpu/janeiro/config-pwr-state.h
+++ b/drivers/edgetpu/janeiro/config-pwr-state.h
@@ -24,6 +24,8 @@ enum edgetpu_pwr_state {
 	TPU_ACTIVE_NOM = 1066000,
 };
 
+#define MIN_ACTIVE_STATE	TPU_ACTIVE_UUD
+
 #define EDGETPU_NUM_STATES 4
 
 extern enum edgetpu_pwr_state edgetpu_active_states[];
diff --git a/drivers/edgetpu/janeiro/config.h b/drivers/edgetpu/janeiro/config.h
index add13dc..a5adbef 100644
--- a/drivers/edgetpu/janeiro/config.h
+++ b/drivers/edgetpu/janeiro/config.h
@@ -13,6 +13,8 @@
 #define EDGETPU_DEV_MAX		1
 
 #define EDGETPU_HAS_MULTI_GROUPS
+/* 1 context per VII/group plus 1 for KCI */
+#define EDGETPU_NCONTEXTS 8
 /* Max number of virtual context IDs that can be allocated for one device. */
 #define EDGETPU_NUM_VCIDS 16
 /* Reserved VCID that uses the extra partition. */
diff --git a/drivers/edgetpu/janeiro/csrs.h b/drivers/edgetpu/janeiro/csrs.h
index ec9598d..385a220 100644
--- a/drivers/edgetpu/janeiro/csrs.h
+++ b/drivers/edgetpu/janeiro/csrs.h
@@ -38,7 +38,6 @@ enum edgetpu_csrs {
 };
 
 /* SYSREG TPU */
-#define EDGETPU_SYSREG_TPU_BASE	0x1cc20000
 #define EDGETPU_SYSREG_TPU_SHAREABILITY	0x700
 #define SHAREABLE_WRITE	(1 << 13)
 #define SHAREABLE_READ	(1 << 12)
diff --git a/drivers/edgetpu/mobile-debug-dump.c b/drivers/edgetpu/mobile-debug-dump.c
index 82dc5a2..3732fbb 100644
--- a/drivers/edgetpu/mobile-debug-dump.c
+++ b/drivers/edgetpu/mobile-debug-dump.c
@@ -7,6 +7,7 @@
  */
 
 #include <linux/mutex.h>
+#include <linux/platform_data/sscoredump.h>
 #include <linux/rbtree.h>
 #include <linux/slab.h>
 
diff --git a/drivers/edgetpu/mobile-debug-dump.h b/drivers/edgetpu/mobile-debug-dump.h
index 105bd70..9eaf069 100644
--- a/drivers/edgetpu/mobile-debug-dump.h
+++ b/drivers/edgetpu/mobile-debug-dump.h
@@ -5,11 +5,10 @@
  *
  * Copyright (C) 2021 Google, Inc.
  */
+
 #ifndef __MOBILE_DEBUG_DUMP_H__
 #define __MOBILE_DEBUG_DUMP_H__
 
-#include <linux/platform_data/sscoredump.h>
-
 #include "edgetpu-debug-dump.h"
 
 struct mobile_sscd_info {
@@ -24,6 +23,8 @@ struct mobile_sscd_mappings_dump {
 	u64 size;
 };
 
+struct sscd_segment;
+
 /*
  * Collects the mapping information of all the host mapping and dmabuf mapping buffers of all
  * @groups as an array of struct mobile_sscd_mappings_dump and populates the @sscd_seg.
diff --git a/drivers/edgetpu/mobile-firmware.c b/drivers/edgetpu/mobile-firmware.c
new file mode 100644
index 0000000..601ecfc
--- /dev/null
+++ b/drivers/edgetpu/mobile-firmware.c
@@ -0,0 +1,490 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Edge TPU firmware management for mobile chipsets.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#include <linux/device.h>
+#include <linux/dma-mapping.h>
+#include <linux/firmware.h>
+#include <linux/gsa/gsa_tpu.h>
+#include <linux/slab.h>
+#include <linux/string.h>
+#include <linux/types.h>
+
+#include "edgetpu.h"
+#include "edgetpu-config.h"
+#include "edgetpu-firmware.h"
+#include "edgetpu-firmware-util.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
+#include "edgetpu-mailbox.h"
+#include "edgetpu-mmu.h"
+#include "edgetpu-mobile-platform.h"
+#include "mobile-firmware.h"
+
+static int mobile_firmware_after_create(struct edgetpu_firmware *et_fw)
+{
+	/*
+	 * Use firmware data to keep a copy of the image config in order
+	 * to avoid re-doing IOMMU mapping on each firmware run
+	 */
+	struct mobile_image_config *data;
+
+	data = kzalloc(sizeof(*data), GFP_KERNEL);
+	if (!data)
+		return -ENOMEM;
+
+	edgetpu_firmware_set_data(et_fw, data);
+	return 0;
+}
+
+static void mobile_firmware_before_destroy(struct edgetpu_firmware *et_fw)
+{
+	struct mobile_image_config *image_config;
+	u32 i, tpu_addr, size;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+
+	image_config = edgetpu_firmware_get_data(et_fw);
+
+	if (image_config && image_config->privilege_level == FW_PRIV_LEVEL_NS) {
+		for (i = 0; i < image_config->num_iommu_mapping; i++) {
+			tpu_addr = image_config->mappings[i].virt_address;
+			if (!tpu_addr)
+				continue;
+			size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
+			edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
+		}
+		edgetpu_firmware_set_data(et_fw, NULL);
+		kfree(image_config);
+	}
+}
+
+static int mobile_firmware_alloc_buffer(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_buffer *fw_buf)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+
+	/* Allocate extra space the image header */
+	size_t buffer_size =
+		etmdev->fw_region_size + MOBILE_FW_HEADER_SIZE;
+
+	fw_buf->vaddr = vmalloc(buffer_size);
+	if (!fw_buf->vaddr) {
+		etdev_err(etdev, "%s: failed to allocate buffer (%zu bytes)\n",
+			  __func__, buffer_size);
+		return -ENOMEM;
+	}
+	fw_buf->dma_addr = 0;
+	fw_buf->alloc_size = buffer_size;
+	fw_buf->used_size_align = 16;
+	return 0;
+}
+
+static void mobile_firmware_free_buffer(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_buffer *fw_buf)
+{
+	vfree(fw_buf->vaddr);
+	fw_buf->vaddr = NULL;
+	fw_buf->dma_addr = 0;
+	fw_buf->alloc_size = 0;
+	fw_buf->used_size_align = 0;
+}
+
+static struct mobile_image_config *mobile_firmware_get_image_config(struct edgetpu_dev *etdev)
+{
+	return (struct mobile_image_config *) edgetpu_firmware_get_data(etdev->firmware);
+}
+
+static void mobile_firmware_save_image_config(struct edgetpu_dev *etdev,
+					      struct mobile_image_config *image_config)
+{
+	struct mobile_image_config *saved_image_config = mobile_firmware_get_image_config(etdev);
+
+	memcpy(saved_image_config, image_config, sizeof(*saved_image_config));
+}
+
+static void mobile_firmware_clear_image_config(struct edgetpu_dev *etdev)
+{
+	struct mobile_image_config *saved_image_config = mobile_firmware_get_image_config(etdev);
+
+	memset(saved_image_config, 0, sizeof(*saved_image_config));
+}
+
+static void mobile_firmware_clear_ns_mappings(struct edgetpu_dev *etdev,
+					      struct mobile_image_config *image_config)
+{
+	unsigned int tpu_addr, size;
+	int i;
+
+	etdev_dbg(etdev, "Firmware image config changed, removing previous mappings\n");
+	for (i = 0; i < image_config->num_iommu_mapping; i++) {
+		tpu_addr = image_config->mappings[i].virt_address;
+		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
+		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
+	}
+}
+
+static int mobile_firmware_setup_ns_mappings(struct edgetpu_firmware *et_fw,
+					     struct mobile_image_config *image_config)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct mobile_image_config *last_image_config = mobile_firmware_get_image_config(etdev);
+	int i, ret;
+	unsigned int tpu_addr, size;
+	phys_addr_t phys_addr;
+
+	if (memcmp(image_config, last_image_config, sizeof(*image_config))) {
+		mobile_firmware_clear_ns_mappings(etdev, last_image_config);
+	} else {
+		return 0;
+	}
+
+	/* Clear saved image config. It will be updated once the new config is confirmed valid */
+	mobile_firmware_clear_image_config(etdev);
+
+	for (i = 0; i < image_config->num_iommu_mapping; i++) {
+		tpu_addr = image_config->mappings[i].virt_address;
+		if (!tpu_addr) {
+			etdev_warn(etdev, "Invalid firmware header\n");
+			goto err;
+		}
+		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
+		phys_addr = (image_config->mappings[i].image_config_value & ~(0xFFF));
+
+		etdev_dbg(etdev, "Adding IOMMU mapping for firmware : %08X -> %08llX", tpu_addr,
+			  phys_addr);
+
+		ret = edgetpu_mmu_add_translation(etdev, tpu_addr, phys_addr, size,
+						  IOMMU_READ | IOMMU_WRITE, EDGETPU_CONTEXT_KCI);
+		if (ret) {
+			etdev_err(etdev,
+				   "Unable to Map: %d tpu_addr: %#x phys_addr: %#llx size: %#x\n",
+				  ret, tpu_addr, phys_addr, size);
+			goto err;
+		}
+	}
+
+	mobile_firmware_save_image_config(etdev, image_config);
+	return 0;
+
+err:
+	while (i--) {
+		tpu_addr = image_config->mappings[i].virt_address;
+		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
+		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
+	}
+	return ret;
+}
+
+static int mobile_firmware_gsa_authenticate(struct edgetpu_mobile_platform_dev *etmdev,
+					    struct edgetpu_firmware_buffer *fw_buf,
+					    struct mobile_image_config *image_config,
+					    void *image_vaddr)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct mobile_image_config *saved_image_config = mobile_firmware_get_image_config(etdev);
+	struct mobile_image_config last_image_config;
+	void *header_vaddr;
+	dma_addr_t header_dma_addr;
+	int tpu_state;
+	int ret = 0;
+	bool image_config_changed = memcmp(image_config, saved_image_config, sizeof(*image_config));
+
+	/* Need a local copy of the saved image in case we need to undo NS mappings */
+	memcpy(&last_image_config, saved_image_config, sizeof(last_image_config));
+
+	/* Clear saved image config. It will be updated once the new config is confirmed valid */
+	mobile_firmware_clear_image_config(etdev);
+
+	tpu_state = gsa_send_tpu_cmd(etmdev->gsa_dev, GSA_TPU_GET_STATE);
+
+	if (tpu_state < GSA_TPU_STATE_INACTIVE) {
+		etdev_err(etdev, "GSA failed to retrieve current status: %d\n", tpu_state);
+		return tpu_state;
+	}
+
+	etdev_dbg(etdev, "GSA Reports TPU state: %d\n", tpu_state);
+
+	if (tpu_state > GSA_TPU_STATE_INACTIVE) {
+		ret = gsa_unload_tpu_fw_image(etmdev->gsa_dev);
+		if (ret) {
+			etdev_warn(etdev, "GSA release failed: %d\n", ret);
+			return -EIO;
+		}
+	}
+
+	/* Copy the firmware image to the target location, skipping the header */
+	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
+	       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
+
+	/* Allocate coherent memory for the image header */
+	header_vaddr = dma_alloc_coherent(etmdev->gsa_dev,
+					  MOBILE_FW_HEADER_SIZE,
+					  &header_dma_addr, GFP_KERNEL);
+	if (!header_vaddr) {
+		etdev_err(etdev,
+			  "Failed to allocate coherent memory for header\n");
+		return -ENOMEM;
+	}
+
+	memcpy(header_vaddr, fw_buf->vaddr, MOBILE_FW_HEADER_SIZE);
+	etdev_dbg(etdev, "Requesting GSA image load. meta = %llX payload = %llX", header_dma_addr,
+		  (u64)etmdev->fw_region_paddr);
+
+	ret = gsa_load_tpu_fw_image(etmdev->gsa_dev, header_dma_addr,
+				    etmdev->fw_region_paddr);
+	if (ret) {
+		etdev_err(etdev, "GSA authentication failed: %d\n", ret);
+	} else {
+		/* Transitioning from NS to something else, clear NS mappings */
+		if (last_image_config.privilege_level == FW_PRIV_LEVEL_NS &&
+		    image_config->privilege_level != FW_PRIV_LEVEL_NS)
+			mobile_firmware_clear_ns_mappings(etdev, &last_image_config);
+
+		/*
+		 * No change in image config, save it here so mobile_firmware_setup_ns_mappings()
+		 * can re-use existing mappings
+		 */
+		if (!image_config_changed)
+			mobile_firmware_save_image_config(etdev, image_config);
+
+		if (image_config->privilege_level == FW_PRIV_LEVEL_NS)
+			ret = mobile_firmware_setup_ns_mappings(etdev->firmware, image_config);
+	}
+
+	dma_free_coherent(etmdev->gsa_dev, MOBILE_FW_HEADER_SIZE, header_vaddr, header_dma_addr);
+
+	return ret;
+}
+
+/* TODO(b/197074886): remove once rio has GSA support */
+static void program_iremap_csr(struct edgetpu_dev *etdev)
+{
+	const int ctx_id = 0, sid0 = 0x30, sid1 = 0x34;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	phys_addr_t fw_paddr = etmdev->fw_region_paddr;
+
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY, (ctx_id << 16) | sid0);
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY + 8,
+			     (ctx_id << 16) | sid1);
+#if defined(ZEBU_SYSMMU_WORKAROUND)
+	/*
+	 * This is required on ZeBu after b/197718405 is fixed, which forwards all transactions to
+	 * the non-secure SysMMU.
+	 */
+	fw_paddr = EDGETPU_INSTRUCTION_REMAP_BASE;
+#endif
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE, fw_paddr);
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE + 8, fw_paddr);
+
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT,
+			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_16M);
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT + 8,
+			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_16M);
+
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL, 1);
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL + 8, 1);
+}
+
+static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
+				       struct edgetpu_firmware_buffer *fw_buf)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+
+	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
+	edgetpu_mailbox_reset(etdev->kci->mailbox);
+
+	if (IS_ENABLED(CONFIG_RIO))
+		program_iremap_csr(etdev);
+
+	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
+}
+
+static int mobile_firmware_restart(struct edgetpu_firmware *et_fw, bool force_reset)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+
+	/*
+	 * We are in a bad state, reset the CPU and hope the device recovers.
+	 * Ignore failures in the reset assert request and proceed to reset release.
+	 */
+	if (force_reset)
+		edgetpu_mobile_firmware_reset_cpu(etdev, true);
+
+	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
+}
+
+static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
+					struct edgetpu_firmware_buffer *fw_buf)
+{
+	int ret = 0;
+	void *image_vaddr;
+	struct mobile_image_config *image_config;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	phys_addr_t image_start, image_end, carveout_start, carveout_end;
+
+	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
+		etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
+			  fw_buf->used_size, MOBILE_FW_HEADER_SIZE);
+		return -EINVAL;
+	}
+
+	image_vaddr = memremap(etmdev->fw_region_paddr,
+			       etmdev->fw_region_size, MEMREMAP_WC);
+	if (!image_vaddr) {
+		etdev_err(etdev, "memremap failed\n");
+		return -ENOMEM;
+	}
+
+	/* fetch the firmware versions */
+	image_config = fw_buf->vaddr + MOBILE_IMAGE_CONFIG_OFFSET;
+	memcpy(&etdev->fw_version, &image_config->firmware_versions,
+	       sizeof(etdev->fw_version));
+
+	if (etmdev->gsa_dev) {
+		ret = mobile_firmware_gsa_authenticate(etmdev, fw_buf, image_config, image_vaddr);
+	} else if (image_config->privilege_level == FW_PRIV_LEVEL_NS) {
+		etdev_dbg(etdev, "Loading unauthenticated non-secure firmware\n");
+		/* Copy the firmware image to the target location, skipping the header */
+		memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
+		       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
+		ret = mobile_firmware_setup_ns_mappings(et_fw, image_config);
+	} else {
+		etdev_err(etdev,
+			  "Cannot load firmware at privilege level %d with no authentication\n",
+			  image_config->privilege_level);
+		ret = -EINVAL;
+	}
+
+	if (ret)
+		goto out;
+
+	image_start = (phys_addr_t)image_config->carveout_base;
+	image_end = (phys_addr_t)(image_config->firmware_base + image_config->firmware_size - 1);
+	carveout_start = etmdev->fw_region_paddr;
+	carveout_end = carveout_start + etmdev->fw_region_size - 1;
+
+	/* Image must fit within the carveout */
+	if (image_start < carveout_start || image_end > carveout_end) {
+		etdev_err(etdev, "Firmware image doesn't fit in carveout\n");
+		etdev_err(etdev, "Image config: %pap - %pap\n", &image_start, &image_end);
+		etdev_err(etdev, "Carveout: %pap - %pap\n", &carveout_start, &carveout_end);
+		ret = -ERANGE;
+	}
+
+out:
+	memunmap(image_vaddr);
+	return ret;
+}
+
+/* Load firmware for chips that use carveout memory for a single chip. */
+int edgetpu_firmware_chip_load_locked(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_desc *fw_desc, const char *name)
+{
+	int ret;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct device *dev = etdev->dev;
+	const struct firmware *fw;
+	size_t aligned_size;
+
+	ret = request_firmware(&fw, name, dev);
+	if (ret) {
+		etdev_dbg(etdev,
+			  "%s: request '%s' failed: %d\n", __func__, name, ret);
+		return ret;
+	}
+
+	aligned_size = ALIGN(fw->size, fw_desc->buf.used_size_align);
+	if (aligned_size > fw_desc->buf.alloc_size) {
+		etdev_dbg(etdev,
+			   "%s: firmware buffer too small: alloc size=%#zx, required size=%#zx\n",
+			  __func__, fw_desc->buf.alloc_size, aligned_size);
+		ret = -ENOSPC;
+		goto out_release_firmware;
+	}
+
+	memcpy(fw_desc->buf.vaddr, fw->data, fw->size);
+	fw_desc->buf.used_size = aligned_size;
+	/* May return NULL on out of memory, driver must handle properly */
+	fw_desc->buf.name = kstrdup(name, GFP_KERNEL);
+
+out_release_firmware:
+	release_firmware(fw);
+	return ret;
+}
+
+void edgetpu_firmware_chip_unload_locked(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_desc *fw_desc)
+{
+	kfree(fw_desc->buf.name);
+	fw_desc->buf.name = NULL;
+	fw_desc->buf.used_size = 0;
+}
+
+int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct mobile_image_config *image_config = mobile_firmware_get_image_config(etdev);
+	int ret = 0;
+
+	if (image_config->privilege_level == FW_PRIV_LEVEL_NS) {
+		int i;
+
+		for (i = 0; i < EDGETPU_NUM_CORES; i++)
+			edgetpu_dev_write_32_sync(etdev, EDGETPU_REG_RESET_CONTROL + i * 8,
+						  assert_reset ? 1 : 0);
+	}
+	else if (etmdev->gsa_dev)
+		ret = gsa_send_tpu_cmd(etmdev->gsa_dev,
+				       assert_reset ? GSA_TPU_SHUTDOWN : GSA_TPU_START);
+	else
+		ret = -ENODEV;
+
+	etdev_dbg(etdev, "%s CPU reset result = %d", assert_reset ? "assert" : "release", ret);
+
+	if (ret < 0)
+		return ret;
+
+	return 0;
+}
+
+static const struct edgetpu_firmware_chip_data mobile_firmware_chip_data = {
+	.default_firmware_name = EDGETPU_DEFAULT_FIRMWARE_NAME,
+	.after_create = mobile_firmware_after_create,
+	.before_destroy = mobile_firmware_before_destroy,
+	.alloc_buffer = mobile_firmware_alloc_buffer,
+	.free_buffer = mobile_firmware_free_buffer,
+	.setup_buffer = mobile_firmware_setup_buffer,
+	.prepare_run = mobile_firmware_prepare_run,
+	.restart = mobile_firmware_restart,
+};
+
+int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev)
+{
+	return edgetpu_firmware_create(etdev, &mobile_firmware_chip_data);
+}
+
+void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev)
+{
+	edgetpu_firmware_destroy(etdev);
+}
+
+unsigned long edgetpu_chip_firmware_iova(struct edgetpu_dev *etdev)
+{
+	/*
+	 * On mobile platforms, firmware address translation may happen in 1 or 2 stages:
+	 * 1.- Instruction remap registers.
+	 * 2.- IOMMU translation (when not running in GSA privilege).
+	 * In either case, the address seen by the TPU's CPU core will remain constant, and
+	 * equal to the macro below.
+	 */
+	return EDGETPU_INSTRUCTION_REMAP_BASE;
+}
diff --git a/drivers/edgetpu/mobile-firmware.h b/drivers/edgetpu/mobile-firmware.h
index 691eaf5..1250d7e 100644
--- a/drivers/edgetpu/mobile-firmware.h
+++ b/drivers/edgetpu/mobile-firmware.h
@@ -12,6 +12,12 @@
 #include "edgetpu-internal.h"
 #include "edgetpu.h"
 
+#define MAX_IOMMU_MAPPINGS 26
+
+#define FW_PRIV_LEVEL_GSA		(0)
+#define FW_PRIV_LEVEL_TZ		(1)
+#define FW_PRIV_LEVEL_NS		(2)
+
 /* mobile FW header size */
 #define MOBILE_FW_HEADER_SIZE SZ_4K
 /* The offset to the signed firmware header. */
@@ -19,6 +25,18 @@
 /* The offset to image configuration. */
 #define MOBILE_IMAGE_CONFIG_OFFSET (MOBILE_HEADER_OFFSET + 0x160)
 
+#define CONFIG_TO_SIZE(a) ((1 << ((a) & 0xFFF)) << 12)
+
+struct iommu_mapping {
+	/* TPU virt address */
+	__u32 virt_address;
+	/*
+	 * contains a 12-bit aligned address and a page-order size into a
+	 * 32-bit value i.e. a physical address and size in page order.
+	 */
+	__u32 image_config_value;
+};
+
 /*
  * The image configuration attached to the signed firmware.
  */
@@ -27,6 +45,12 @@ struct mobile_image_config {
 	__u32 firmware_base;
 	__u32 firmware_size;
 	struct edgetpu_fw_version firmware_versions;
+	__u32 config_version;
+	__u32 privilege_level;
+	__u32 remapped_region_start;
+	__u32 remapped_region_end;
+	__u32 num_iommu_mapping;
+	struct iommu_mapping mappings[MAX_IOMMU_MAPPINGS];
 } __packed;
 
 /*
@@ -46,7 +70,15 @@ struct mobile_image_header {
 	struct mobile_image_config ImageConfig;
 };
 
-int mobile_edgetpu_firmware_create(struct edgetpu_dev *etdev);
-void mobile_edgetpu_firmware_destroy(struct edgetpu_dev *etdev);
+int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev);
+void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev);
+
+/*
+ * Assert or release the reset signal of the TPU's CPU
+ * Depending on privilege level, this may be by a direct register write
+ * or a call into GSA.
+ */
+int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset);
+
 
 #endif /* __MOBILE_FIRMWARE_H__ */
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
new file mode 100644
index 0000000..de21bda
--- /dev/null
+++ b/drivers/edgetpu/mobile-pm.c
@@ -0,0 +1,570 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Common EdgeTPU mobile power management support
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#include <linux/atomic.h>
+#include <linux/delay.h>
+#include <linux/gsa/gsa_tpu.h>
+#include <linux/module.h>
+#include <linux/pm_runtime.h>
+#include <soc/google/bcl.h>
+#include <soc/google/bts.h>
+#include <soc/google/exynos_pm_qos.h>
+
+#include "edgetpu-config.h"
+#include "edgetpu-firmware.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
+#include "edgetpu-mailbox.h"
+#include "edgetpu-mobile-platform.h"
+#include "edgetpu-pm.h"
+#include "mobile-firmware.h"
+#include "mobile-pm.h"
+
+#include "edgetpu-pm.c"
+
+/*
+ * Encode INT/MIF values as a 16 bit pair in the 32-bit return value
+ * (in units of MHz, to provide enough range)
+ */
+#define PM_QOS_INT_SHIFT                (16)
+#define PM_QOS_MIF_MASK                 (0xFFFF)
+#define PM_QOS_FACTOR                   (1000)
+
+static int power_state = TPU_DEFAULT_POWER_STATE;
+
+module_param(power_state, int, 0660);
+
+#define MAX_VOLTAGE_VAL 1250000
+
+enum edgetpu_pwr_state edgetpu_active_states[EDGETPU_NUM_STATES] = {
+	TPU_ACTIVE_UUD,
+	TPU_ACTIVE_SUD,
+	TPU_ACTIVE_UD,
+	TPU_ACTIVE_NOM,
+#if IS_ENABLED(CONFIG_ABROLHOS)
+	TPU_ACTIVE_OD,
+#endif /* IS_ENABLED(CONFIG_ABROLHOS) */
+};
+
+uint32_t *edgetpu_states_display = edgetpu_active_states;
+
+static int mobile_pwr_state_init(struct device *dev)
+{
+	int ret;
+	int curr_state;
+
+	pm_runtime_enable(dev);
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+
+	if (curr_state > TPU_OFF) {
+		ret = pm_runtime_get_sync(dev);
+		if (ret) {
+			dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
+			return ret;
+		}
+	}
+
+	ret = exynos_acpm_set_init_freq(TPU_ACPM_DOMAIN, curr_state);
+	if (ret) {
+		dev_err(dev, "error initializing tpu state: %d\n", ret);
+		if (curr_state > TPU_OFF)
+			pm_runtime_put_sync(dev);
+		return ret;
+	}
+
+	return ret;
+}
+
+static int mobile_pwr_state_set_locked(struct edgetpu_mobile_platform_dev *etmdev, u64 val)
+{
+	int ret;
+	int curr_state;
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	struct device *dev = etdev->dev;
+
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+
+	dev_dbg(dev, "Power state %d -> %llu\n", curr_state, val);
+
+	if (curr_state == TPU_OFF && val > TPU_OFF) {
+		ret = pm_runtime_get_sync(dev);
+		if (ret) {
+			dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
+			return ret;
+		}
+	}
+
+	ret = platform_pwr->acpm_set_rate(TPU_ACPM_DOMAIN, (unsigned long)val);
+	if (ret) {
+		dev_err(dev, "error setting tpu state: %d\n", ret);
+		pm_runtime_put_sync(dev);
+		return ret;
+	}
+
+	if (curr_state != TPU_OFF && val == TPU_OFF) {
+		ret = pm_runtime_put_sync(dev);
+		if (ret) {
+			dev_err(dev, "%s: pm_runtime_put_sync returned %d\n", __func__, ret);
+			return ret;
+		}
+		if (platform_pwr->block_down)
+			platform_pwr->block_down(etdev);
+	}
+
+	return ret;
+}
+
+static int mobile_pwr_state_get_locked(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct device *dev = etdev->dev;
+
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	dev_dbg(dev, "current tpu state: %llu\n", *val);
+
+	return 0;
+}
+
+static int mobile_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->requested_state = val;
+	if (val >= platform_pwr->min_state)
+		ret = mobile_pwr_state_set_locked(etmdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int mobile_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int ret;
+
+	mutex_lock(&platform_pwr->state_lock);
+	ret = mobile_pwr_state_get_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int mobile_min_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->min_state = val;
+	if (val >= platform_pwr->requested_state)
+		ret = mobile_pwr_state_set_locked(etmdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int mobile_min_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	mutex_lock(&platform_pwr->state_lock);
+	*val = platform_pwr->min_state;
+	mutex_unlock(&platform_pwr->state_lock);
+	return 0;
+}
+
+static int mobile_pwr_policy_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int ret;
+
+	mutex_lock(&platform_pwr->policy_lock);
+	ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, val);
+
+	if (ret) {
+		dev_err(etmdev->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 mobile_pwr_policy_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	mutex_lock(&platform_pwr->policy_lock);
+	*val = platform_pwr->curr_policy;
+	mutex_unlock(&platform_pwr->policy_lock);
+
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_policy, mobile_pwr_policy_get, mobile_pwr_policy_set,
+			 "%llu\n");
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, mobile_pwr_state_get, mobile_pwr_state_set, "%llu\n");
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, mobile_min_pwr_state_get, mobile_min_pwr_state_set,
+			 "%llu\n");
+
+static int mobile_get_initial_pwr_state(struct device *dev)
+{
+	switch (power_state) {
+#if IS_ENABLED(CONFIG_ABROLHOS)
+	case TPU_ACTIVE_OD:
+	case TPU_DEEP_SLEEP_CLOCKS_SLOW:
+	case TPU_DEEP_SLEEP_CLOCKS_FAST:
+	case TPU_RETENTION_CLOCKS_SLOW:
+#endif /* IS_ENABLED(CONFIG_ABROLHOS) */
+	case TPU_ACTIVE_UUD:
+	case TPU_ACTIVE_SUD:
+	case TPU_ACTIVE_UD:
+	case TPU_ACTIVE_NOM:
+		dev_info(dev, "Initial power state: %d\n", power_state);
+		break;
+	case TPU_OFF:
+#if IS_ENABLED(CONFIG_ABROLHOS)
+	case TPU_DEEP_SLEEP_CLOCKS_OFF:
+	case TPU_SLEEP_CLOCKS_OFF:
+#endif /* IS_ENABLED(CONFIG_ABROLHOS) */
+		dev_warn(dev, "Power state %d prevents control core booting", power_state);
+		fallthrough;
+	default:
+		dev_warn(dev, "Power state %d is invalid\n", power_state);
+		dev_warn(dev, "defaulting to active nominal\n");
+		power_state = TPU_ACTIVE_NOM;
+		break;
+	}
+	return power_state;
+}
+
+static void mobile_power_down(struct edgetpu_pm *etpm);
+
+static int mobile_power_up(struct edgetpu_pm *etpm)
+{
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int ret = mobile_pwr_state_set(etpm->etdev, mobile_get_initial_pwr_state(etdev->dev));
+
+	etdev_info(etpm->etdev, "Powering up\n");
+
+	if (ret)
+		return ret;
+
+	if (platform_pwr->lpm_up)
+		platform_pwr->lpm_up(etdev);
+
+	edgetpu_chip_init(etdev);
+
+	if (etdev->kci) {
+		etdev_dbg(etdev, "Resetting KCI\n");
+		edgetpu_kci_reinit(etdev->kci);
+	}
+	if (etdev->mailbox_manager) {
+		etdev_dbg(etdev, "Resetting VII mailboxes\n");
+		edgetpu_mailbox_reset_vii(etdev->mailbox_manager);
+	}
+
+	if (!etdev->firmware)
+		return 0;
+
+	/*
+	 * Why this function uses edgetpu_firmware_*_locked functions without explicitly holding
+	 * edgetpu_firmware_lock:
+	 *
+	 * edgetpu_pm_get() is called in two scenarios - one is when the firmware loading is
+	 * attempting, another one is when the user-space clients need the device be powered
+	 * (usually through acquiring the wakelock).
+	 *
+	 * For the first scenario edgetpu_firmware_is_loading() below shall return true.
+	 * For the second scenario we are indeed called without holding the firmware lock, but the
+	 * firmware loading procedures (i.e. the first scenario) always call edgetpu_pm_get() before
+	 * changing the firmware state, and edgetpu_pm_get() is blocked until this function
+	 * finishes. In short, we are protected by the PM lock.
+	 */
+
+	if (edgetpu_firmware_is_loading(etdev))
+		return 0;
+
+	/* attempt firmware run */
+	switch (edgetpu_firmware_status_locked(etdev)) {
+	case FW_VALID:
+		ret = edgetpu_firmware_restart_locked(etdev, false);
+		break;
+	case FW_INVALID:
+		ret = edgetpu_firmware_run_default_locked(etdev);
+		break;
+	default:
+		break;
+	}
+	if (ret) {
+		mobile_power_down(etpm);
+	} else {
+#if IS_ENABLED(CONFIG_GOOGLE_BCL)
+		if (!etmdev->bcl_dev)
+			etmdev->bcl_dev = google_retrieve_bcl_handle();
+		if (etmdev->bcl_dev)
+			google_init_tpu_ratio(etmdev->bcl_dev);
+#endif
+	}
+
+	return ret;
+}
+
+static void mobile_pm_cleanup_bts_scenario(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	if (!performance_scenario)
+		return;
+
+	mutex_lock(&platform_pwr->scenario_lock);
+	while (platform_pwr->scenario_count) {
+		int ret = bts_del_scenario(performance_scenario);
+
+		if (ret) {
+			platform_pwr->scenario_count = 0;
+			etdev_warn_once(etdev, "error %d in cleaning up BTS scenario %u\n", ret,
+					performance_scenario);
+			break;
+		}
+		platform_pwr->scenario_count--;
+	}
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+static void mobile_power_down(struct edgetpu_pm *etpm)
+{
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	u64 val;
+	int res = 0;
+	int min_state = platform_pwr->min_state;
+
+	etdev_info(etdev, "Powering down\n");
+
+	if (min_state >= MIN_ACTIVE_STATE) {
+		etdev_info(etdev, "Power down skipped due to min state = %d\n", min_state);
+		return;
+	}
+
+	if (mobile_pwr_state_get(etdev, &val)) {
+		etdev_warn(etdev, "Failed to read current power state\n");
+		val = TPU_ACTIVE_NOM;
+	}
+	if (val == TPU_OFF) {
+		etdev_dbg(etdev, "Device already off, skipping shutdown\n");
+		return;
+	}
+
+	if (edgetpu_firmware_status_locked(etdev) == FW_VALID) {
+		/* Update usage stats before we power off fw. */
+		edgetpu_kci_update_usage_locked(etdev);
+		platform_pwr->firmware_down(etdev);
+		edgetpu_kci_cancel_work_queues(etdev->kci);
+		res = edgetpu_mobile_firmware_reset_cpu(etdev, true);
+		if (res < 0)
+			etdev_warn(etdev, "CPU reset request failed (%d)\n", res);
+	}
+
+	if (platform_pwr->lpm_down)
+		platform_pwr->lpm_down(etdev);
+
+	mobile_pwr_state_set(etdev, TPU_OFF);
+
+	/* Remove our vote for INT/MIF state (if any) */
+	exynos_pm_qos_update_request(&platform_pwr->int_min, 0);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, 0);
+
+	mobile_pm_cleanup_bts_scenario(etdev);
+
+	/*
+	 * It should be impossible that power_down() is called when secure_client is set.
+	 * Non-null secure_client implies ext mailbox is acquired, which implies wakelock is
+	 * acquired.
+	 * Clear the state here just in case.
+	 */
+	etmdev->secure_client = NULL;
+}
+
+static int mobile_pm_after_create(struct edgetpu_pm *etpm)
+{
+	int ret;
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct device *dev = etdev->dev;
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	ret = mobile_pwr_state_init(dev);
+	if (ret)
+		return ret;
+
+	mutex_init(&platform_pwr->policy_lock);
+	mutex_init(&platform_pwr->state_lock);
+	mutex_init(&platform_pwr->scenario_lock);
+
+	exynos_pm_qos_add_request(&platform_pwr->int_min, PM_QOS_DEVICE_THROUGHPUT, 0);
+	exynos_pm_qos_add_request(&platform_pwr->mif_min, PM_QOS_BUS_THROUGHPUT, 0);
+
+	platform_pwr->performance_scenario = bts_get_scenindex("tpu_performance");
+	if (!platform_pwr->performance_scenario)
+		dev_warn(etdev->dev, "tpu_performance BTS scenario not found\n");
+	platform_pwr->scenario_count = 0;
+
+	ret = mobile_pwr_state_set(etdev, mobile_get_initial_pwr_state(dev));
+	if (ret)
+		return ret;
+	platform_pwr->debugfs_dir = debugfs_create_dir("power", edgetpu_fs_debugfs_dir());
+	if (IS_ERR_OR_NULL(platform_pwr->debugfs_dir)) {
+		dev_warn(etdev->dev, "Failed to create debug FS power");
+		/* don't fail the procedure on debug FS creation fails */
+		return 0;
+	}
+	debugfs_create_file("state", 0660, platform_pwr->debugfs_dir, etdev, &fops_tpu_pwr_state);
+	debugfs_create_file("min_state", 0660, platform_pwr->debugfs_dir, etdev,
+			    &fops_tpu_min_pwr_state);
+	debugfs_create_file("policy", 0660, platform_pwr->debugfs_dir, etdev, &fops_tpu_pwr_policy);
+
+	if (platform_pwr->after_create)
+		ret = platform_pwr->after_create(etdev);
+
+	return ret;
+}
+
+static void mobile_pm_before_destroy(struct edgetpu_pm *etpm)
+{
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	if (platform_pwr->before_destroy)
+		platform_pwr->before_destroy(etdev);
+
+	debugfs_remove_recursive(platform_pwr->debugfs_dir);
+	pm_runtime_disable(etpm->etdev->dev);
+	mobile_pm_cleanup_bts_scenario(etdev);
+	exynos_pm_qos_remove_request(&platform_pwr->int_min);
+	exynos_pm_qos_remove_request(&platform_pwr->mif_min);
+}
+
+static struct edgetpu_pm_handlers mobile_pm_handlers = {
+	.after_create = mobile_pm_after_create,
+	.before_destroy = mobile_pm_before_destroy,
+	.power_up = mobile_power_up,
+	.power_down = mobile_power_down,
+};
+
+int mobile_pm_create(struct edgetpu_dev *etdev)
+{
+	return edgetpu_pm_create(etdev, &mobile_pm_handlers);
+}
+
+void mobile_pm_destroy(struct edgetpu_dev *etdev)
+{
+	edgetpu_pm_destroy(etdev);
+}
+
+void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	s32 int_val = (pm_qos_val >> PM_QOS_INT_SHIFT) * PM_QOS_FACTOR;
+	s32 mif_val = (pm_qos_val & PM_QOS_MIF_MASK) * PM_QOS_FACTOR;
+
+	etdev_dbg(etdev, "%s: pm_qos request - int = %d mif = %d\n", __func__, int_val, mif_val);
+
+	exynos_pm_qos_update_request(&platform_pwr->int_min, int_val);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, mif_val);
+}
+
+static void mobile_pm_activate_bts_scenario(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	/* bts_add_scenario() keeps track of reference count internally.*/
+	int ret;
+
+	if (!performance_scenario)
+		return;
+	mutex_lock(&platform_pwr->scenario_lock);
+	ret = bts_add_scenario(performance_scenario);
+	if (ret)
+		etdev_warn_once(etdev, "error %d adding BTS scenario %u\n", ret,
+				performance_scenario);
+	else
+		platform_pwr->scenario_count++;
+
+	etdev_dbg(etdev, "BTS Scenario activated: %d\n", platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+static void mobile_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
+{
+	/* bts_del_scenario() keeps track of reference count internally.*/
+	int ret;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	if (!performance_scenario)
+		return;
+	mutex_lock(&platform_pwr->scenario_lock);
+	if (!platform_pwr->scenario_count) {
+		etdev_warn(etdev, "Unbalanced bts deactivate\n");
+		mutex_unlock(&platform_pwr->scenario_lock);
+		return;
+	}
+	ret = bts_del_scenario(performance_scenario);
+	if (ret)
+		etdev_warn_once(etdev, "error %d deleting BTS scenario %u\n", ret,
+				performance_scenario);
+	else
+		platform_pwr->scenario_count--;
+
+	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n", platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+void mobile_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val)
+{
+	etdev_dbg(etdev, "%s: bts request - val = %u\n", __func__, bts_val);
+
+	switch (bts_val) {
+	case 0:
+		mobile_pm_deactivate_bts_scenario(etdev);
+		break;
+	case 1:
+		mobile_pm_activate_bts_scenario(etdev);
+		break;
+	default:
+		etdev_warn(etdev, "%s: invalid BTS request value: %u\n", __func__, bts_val);
+		break;
+	}
+}
diff --git a/drivers/edgetpu/mobile-pm.h b/drivers/edgetpu/mobile-pm.h
index a04b6bb..e5fafb6 100644
--- a/drivers/edgetpu/mobile-pm.h
+++ b/drivers/edgetpu/mobile-pm.h
@@ -52,10 +52,17 @@ enum mobile_reverse_kci_code {
 /*
  * Initialize a power management interface for an edgetpu device on mobile
  * chipsets.
+ * Needs to be called after the devices's platform_pwr struct has been initialized.
  */
 int mobile_pm_create(struct edgetpu_dev *etdev);
 
 /*
+ * Wrapper for chip-specific implementation.
+ * Typically calls mobile_pm_create after initializing the platform_pwr struct.
+ */
+int edgetpu_chip_pm_create(struct edgetpu_dev *etdev);
+
+/*
  * Destroy power management interface for an edgetpu device on mobile chipsets.
  */
 void mobile_pm_destroy(struct edgetpu_dev *etdev);