Merge android13-gs-pixel-5.10-gs101 into android13-gs-pixel-5.10

SBMerger: 410055097
Change-Id: Ie5c7ea49e6441bff8df922f7959b8750517500fc
Signed-off-by: SecurityBot <android-nexus-securitybot@system.gserviceaccount.com>

14 files changed, 312 insertions, 154 deletions

diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index 67feb70..0aecdd1 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -99,6 +99,13 @@ void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
 	case RKCI_CODE_BTS:
 		mobile_pm_set_bts(etdev, resp->retval);
 		break;
+	case RKCI_CODE_PM_QOS_BTS:
+		/* FW indicates to ignore the request by setting them to undefined values. */
+		if (resp->retval != (typeof(resp->retval))~0ull)
+			mobile_pm_set_pm_qos(etdev, resp->retval);
+		if (resp->status != (typeof(resp->status))~0ull)
+			mobile_pm_set_bts(etdev, resp->status);
+		break;
 	default:
 		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n",
 			   __func__, resp->code);
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index deb763f..710adbf 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -1,2 +1,4 @@
 // SPDX-License-Identifier: GPL-2.0
+#include <linux/thermal.h>
+#include <soc/google/gs101_tmu.h>
 #include "mobile-thermal.c"
diff --git a/drivers/edgetpu/edgetpu-debug-dump.c b/drivers/edgetpu/edgetpu-debug-dump.c
index 86928a6..a017e7f 100644
--- a/drivers/edgetpu/edgetpu-debug-dump.c
+++ b/drivers/edgetpu/edgetpu-debug-dump.c
@@ -73,7 +73,10 @@ int edgetpu_get_debug_dump(struct edgetpu_dev *etdev, u64 type)
 		if (init_fw_dump_buffer)
 			etdev_err(etdev, "failed to init dump buffer in FW");
 
-		etdev_err(etdev, "Debug dump KCI req failed: %d", ret);
+		if (ret == KCI_ERROR_UNIMPLEMENTED)
+			etdev_dbg(etdev, "Debug dump KCI not implemented");
+		else
+			etdev_err(etdev, "Debug dump KCI req failed: %d", ret);
 	}
 
 	return ret;
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 64fbbfa..49c672d 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -116,6 +116,14 @@ static int edgetpu_kci_leave_group_worker(struct kci_worker_param *param)
 
 #endif /* EDGETPU_HAS_MCP */
 
+static int edgetpu_group_activate_external_mailbox(struct edgetpu_device_group *group)
+{
+	if (!group->ext_mailbox)
+		return 0;
+	edgetpu_mailbox_reinit_external_mailbox(group);
+	return edgetpu_mailbox_activate_external_mailbox(group);
+}
+
 /*
  * Activates the VII mailbox @group owns.
  *
@@ -146,6 +154,12 @@ static int edgetpu_group_activate(struct edgetpu_device_group *group)
 	return ret;
 }
 
+static void edgetpu_group_deactivate_external_mailbox(struct edgetpu_device_group *group)
+{
+	edgetpu_mailbox_deactivate_external_mailbox(group);
+	edgetpu_mailbox_disable_external_mailbox(group);
+}
+
 /*
  * Deactivates the VII mailbox @group owns.
  *
@@ -1166,8 +1180,8 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	 */
 	pages = kvmalloc((num_pages * sizeof(*pages)), GFP_KERNEL | __GFP_NOWARN);
 	if (!pages) {
-		etdev_dbg(etdev, "%s: kvmalloc pages failed (%lu bytes)\n",
-			  __func__, (num_pages * sizeof(*pages)));
+		etdev_err(etdev, "out of memory allocating pages (%lu bytes)",
+			  num_pages * sizeof(*pages));
 		return ERR_PTR(-ENOMEM);
 	}
 	/*
@@ -1210,8 +1224,8 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	/* Allocate our own vmas array non-contiguous. */
 	vmas = kvmalloc((num_pages * sizeof(*vmas)), GFP_KERNEL | __GFP_NOWARN);
 	if (!vmas) {
-		etdev_dbg(etdev, "%s: kvmalloc vmas failed (%lu bytes)\n",
-			  __func__, (num_pages * sizeof(*pages)));
+		etdev_err(etdev, "out of memory allocating vmas (%lu bytes)",
+			  num_pages * sizeof(*pages));
 		kvfree(pages);
 		return ERR_PTR(-ENOMEM);
 	}
@@ -1223,6 +1237,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 			  group->workload_id, (void *)host_addr, num_pages,
 			  ret);
 		num_pages = 0;
+
+		if (ret == -ENOMEM)
+			etdev_err(etdev,
+				  "system out of memory locking %u pages",
+				  num_pages);
 		goto error;
 	}
 	if (ret < num_pages) {
@@ -1856,6 +1875,7 @@ void edgetpu_group_close_and_detach_mailbox(struct edgetpu_device_group *group)
 	if (is_finalized_or_errored(group)) {
 		edgetpu_group_deactivate(group);
 		edgetpu_group_detach_mailbox_locked(group);
+		edgetpu_group_deactivate_external_mailbox(group);
 	}
 	mutex_unlock(&group->lock);
 }
@@ -1885,8 +1905,14 @@ int edgetpu_group_attach_and_open_mailbox(struct edgetpu_device_group *group)
 		goto out_unlock;
 	ret = edgetpu_group_activate(group);
 	if (ret)
-		edgetpu_group_detach_mailbox_locked(group);
+		goto error_detach;
+	ret = edgetpu_group_activate_external_mailbox(group);
+	if (!ret)
+		goto out_unlock;
 
+	edgetpu_group_deactivate(group);
+error_detach:
+	edgetpu_group_detach_mailbox_locked(group);
 out_unlock:
 	mutex_unlock(&group->lock);
 	return ret;
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 05bd9c9..e722738 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -112,25 +112,20 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 	etdev = client->etdev;
 
 	LOCK(client);
-	wakelock_count = edgetpu_wakelock_lock(client->wakelock);
 	/*
-	 * @wakelock_count = 0 means the device might be powered off. And for group with a
-	 * non-detachable mailbox, its mailbox is removed when the group is released, in such case
-	 * we need to ensure the device is powered to prevent kernel panic on programming VII
-	 * mailbox CSRs.
-	 *
-	 * For mailbox-detachable groups the mailbox had been removed when the wakelock was
-	 * released, edgetpu_device_group_release() doesn't need the device be powered in this case.
+	 * Safe to read wakelock->req_count here since req_count is only
+	 * modified during [acquire/release]_wakelock ioctl calls which cannot
+	 * race with releasing client/fd.
 	 */
-	if (!wakelock_count && client->group && !client->group->mailbox_detachable) {
-		/* assumes @group->etdev == @client->etdev, i.e. @client is the leader of @group */
-		if (!edgetpu_pm_get(etdev->pm))
-			wakelock_count = 1;
-		else
-			/* failed to power on - prevent group releasing from accessing the device */
-			client->group->dev_inaccessible = true;
-	}
-	edgetpu_wakelock_unlock(client->wakelock);
+	wakelock_count = NO_WAKELOCK(client->wakelock) ? 1 : client->wakelock->req_count;
+	/*
+	 * @wakelock_count = 0 means the device might be powered off. Mailbox(EXT/VII) is removed
+	 * when the group is released, So we need to ensure the device should not accessed to
+	 * prevent kernel panic on programming mailbox CSRs.
+	 */
+	if (!wakelock_count && client->group)
+		client->group->dev_inaccessible = true;
+
 	UNLOCK(client);
 
 	edgetpu_client_remove(client);
@@ -612,15 +607,20 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 		ret = count;
 		goto error_unlock;
 	}
-	if (!count && client->group)
-		ret = edgetpu_group_attach_and_open_mailbox(client->group);
-	if (ret) {
-		etdev_warn(client->etdev,
-			   "failed to attach mailbox: %d", ret);
+	if (!count) {
+		if (client->group)
+			ret = edgetpu_group_attach_and_open_mailbox(client->group);
+		if (ret) {
+			etdev_warn(client->etdev,
+				   "failed to attach mailbox: %d", ret);
+			edgetpu_pm_put(client->etdev->pm);
+			edgetpu_wakelock_release(client->wakelock);
+			edgetpu_wakelock_unlock(client->wakelock);
+			goto error_unlock;
+		}
+	} else {
+		/* Balance the power up count due to pm_get above.*/
 		edgetpu_pm_put(client->etdev->pm);
-		edgetpu_wakelock_release(client->wakelock);
-		edgetpu_wakelock_unlock(client->wakelock);
-		goto error_unlock;
 	}
 	edgetpu_wakelock_unlock(client->wakelock);
 	UNLOCK(client);
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index 0bd68b8..0f28918 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -354,12 +354,15 @@ int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
 	 * Per-context mappings are mirrored to their specific domains here
 	 */
 	if (params.domain != default_domain) {
-		if (!iommu_map_sg(params.domain, iova, map->sgt.sgl,
-				  map->sgt.orig_nents, params.prot)) {
+		ssize_t mapped = (ssize_t)iommu_map_sg(params.domain, iova, map->sgt.sgl,
+						       map->sgt.orig_nents, params.prot);
+
+		/* iommu_map_sg returns 0 on failure before 5.15, returns -errno afterwards */
+		if (mapped <= 0) {
 			/* Undo the mapping in the default domain */
 			dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.orig_nents, map->dir,
 					   DMA_ATTR_SKIP_CPU_SYNC);
-			return -ENOMEM;
+			return mapped == 0 ? -ENOMEM : (int)mapped;
 		}
 	}
 
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 138242e..834318f 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -240,7 +240,7 @@ static void edgetpu_kci_consume_wait_list(
  */
 static void
 edgetpu_kci_handle_response(struct edgetpu_kci *kci,
-			    const struct edgetpu_kci_response_element *resp)
+			    struct edgetpu_kci_response_element *resp)
 {
 	if (resp->seq & KCI_REVERSE_FLAG) {
 		int ret = edgetpu_reverse_kci_add_response(kci, resp);
@@ -252,6 +252,11 @@ edgetpu_kci_handle_response(struct edgetpu_kci *kci,
 				resp->code, ret);
 		return;
 	}
+	/*
+	 * Response status can be set after the RKCI status has been preserved
+	 * by copying response elements in the circular queue above.
+	 */
+	resp->status = KCI_STATUS_OK;
 	edgetpu_kci_consume_wait_list(kci, resp);
 }
 
@@ -322,7 +327,6 @@ static struct edgetpu_kci_response_element *edgetpu_kci_fetch_responses(
 		j = CIRCULAR_QUEUE_REAL_INDEX(head);
 		for (i = 0; i < count; i++) {
 			memcpy(&ret[total], &queue[j], sizeof(*queue));
-			ret[total].status = KCI_STATUS_OK;
 			j = (j + 1) % size;
 			total++;
 		}
@@ -400,7 +404,6 @@ edgetpu_kci_fetch_one_response(struct edgetpu_kci *kci,
 	}
 
 	memcpy(resp, &queue[CIRCULAR_QUEUE_REAL_INDEX(head)], sizeof(*queue));
-	resp->status = KCI_STATUS_OK;
 	edgetpu_mailbox_inc_resp_queue_head(kci->mailbox, 1);
 
 	spin_unlock(&kci->resp_queue_lock);
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 007a512..5fac2a2 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -491,6 +491,18 @@ void edgetpu_mailbox_remove_vii(struct edgetpu_vii *vii)
 	}
 }
 
+static int edgetpu_mailbox_do_alloc_queue(struct edgetpu_dev *etdev,
+					  struct edgetpu_mailbox *mailbox, u32 queue_size,
+					  u32 unit, edgetpu_queue_mem *mem)
+{
+	u32 size = unit * queue_size;
+
+	/* Align queue size to page size for TPU MMU map. */
+	size = __ALIGN_KERNEL(size, PAGE_SIZE);
+	return edgetpu_iremap_alloc(etdev, size, mem,
+				    edgetpu_mailbox_context_id(mailbox));
+}
+
 /*
  * Allocates memory for a queue.
  *
@@ -505,21 +517,14 @@ int edgetpu_mailbox_alloc_queue(struct edgetpu_dev *etdev,
 				u32 unit, enum mailbox_queue_type type,
 				edgetpu_queue_mem *mem)
 {
-	u32 size = unit * queue_size;
-	int ret;
-
-	/* Align queue size to page size for TPU MMU map. */
-	size = __ALIGN_KERNEL(size, PAGE_SIZE);
-	ret = edgetpu_iremap_alloc(etdev, size, mem,
-				   edgetpu_mailbox_context_id(mailbox));
+	int ret = edgetpu_mailbox_do_alloc_queue(etdev, mailbox, queue_size, unit, mem);
 	if (ret)
 		return ret;
 
 	ret = edgetpu_mailbox_set_queue(mailbox, type, mem->tpu_addr,
 					queue_size);
 	if (ret) {
-		edgetpu_iremap_free(etdev, mem,
-				    edgetpu_mailbox_context_id(mailbox));
+		edgetpu_mailbox_free_queue(etdev, mailbox, mem);
 		return ret;
 	}
 	return 0;
@@ -660,13 +665,17 @@ void edgetpu_mailbox_init_doorbells(struct edgetpu_mailbox *mailbox)
 	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, resp_queue_doorbell_enable, 1);
 }
 
-void edgetpu_mailbox_reset_vii(struct edgetpu_mailbox_manager *mgr)
+void edgetpu_mailbox_reset_mailboxes(struct edgetpu_mailbox_manager *mgr)
 {
 	uint i;
 	unsigned long flags;
 
 	write_lock_irqsave(&mgr->mailboxes_lock, flags);
-	for (i = mgr->vii_index_from; i < mgr->vii_index_to; i++) {
+	/*
+	 * Reset all the allocated mailboxes, starting from VII till
+	 * external mailboxes.
+	 */
+	for (i = mgr->vii_index_from; i < mgr->ext_index_to; i++) {
 		struct edgetpu_mailbox *mbox = mgr->mailboxes[i];
 
 		if (!mbox)
@@ -712,7 +721,43 @@ void edgetpu_mailbox_reinit_vii(struct edgetpu_device_group *group)
 	edgetpu_mailbox_enable(mailbox);
 }
 
-void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev)
+static void edgetpu_mailbox_init_external_mailbox(struct edgetpu_external_mailbox *ext_mailbox)
+{
+	struct edgetpu_mailbox_attr attr;
+	struct edgetpu_mailbox *mailbox;
+	struct edgetpu_mailbox_descriptor *desc;
+	uint i;
+
+	attr = ext_mailbox->attr;
+
+	for (i = 0; i < ext_mailbox->count; i++) {
+		desc = &ext_mailbox->descriptors[i];
+		mailbox = desc->mailbox;
+		edgetpu_mailbox_set_priority(mailbox, attr.priority);
+		EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, cmd_queue_tail_doorbell_enable,
+					      attr.cmdq_tail_doorbell);
+		edgetpu_mailbox_set_queue(mailbox, MAILBOX_CMD_QUEUE,
+					  desc->cmd_queue_mem.tpu_addr,
+					  attr.cmd_queue_size);
+		edgetpu_mailbox_set_queue(mailbox, MAILBOX_RESP_QUEUE,
+					  desc->resp_queue_mem.tpu_addr,
+					  attr.resp_queue_size);
+		edgetpu_mailbox_enable(mailbox);
+	}
+}
+
+void edgetpu_mailbox_reinit_external_mailbox(struct edgetpu_device_group *group)
+{
+	struct edgetpu_external_mailbox *ext_mailbox = group->ext_mailbox;
+
+	if (!ext_mailbox)
+		return;
+
+	etdev_dbg(group->etdev, "Restoring external attached %d mailboxes\n", ext_mailbox->count);
+	edgetpu_mailbox_init_external_mailbox(ext_mailbox);
+}
+
+void edgetpu_mailbox_restore_active_mailbox_queues(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_list_group *l;
 	struct edgetpu_device_group *group;
@@ -770,8 +815,10 @@ void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev)
 		 * If the group is just finalized or has mailbox attached in
 		 * another process, this re-init is redundant but isn't harmful.
 		 */
-		if (edgetpu_group_finalized_and_attached(group))
+		if (edgetpu_group_finalized_and_attached(group)) {
 			edgetpu_mailbox_reinit_vii(group);
+			edgetpu_mailbox_reinit_external_mailbox(group);
+		}
 		mutex_unlock(&group->lock);
 		edgetpu_device_group_put(group);
 	}
@@ -784,23 +831,22 @@ static int edgetpu_mailbox_external_alloc_queue_batch(struct edgetpu_external_ma
 	struct edgetpu_mailbox *mailbox;
 	struct edgetpu_mailbox_attr attr;
 	struct edgetpu_mailbox_descriptor *desc;
+	struct edgetpu_dev *etdev = ext_mailbox->etdev;
 
 	attr = ext_mailbox->attr;
 
 	for (i = 0; i < ext_mailbox->count; i++) {
 		desc = &ext_mailbox->descriptors[i];
 		mailbox = desc->mailbox;
-		ret = edgetpu_mailbox_alloc_queue(ext_mailbox->etdev, mailbox, attr.cmd_queue_size,
-						  attr.sizeof_cmd, MAILBOX_CMD_QUEUE,
-						  &desc->cmd_queue_mem);
+		ret = edgetpu_mailbox_do_alloc_queue(etdev, mailbox, attr.cmd_queue_size,
+						     attr.sizeof_cmd, &desc->cmd_queue_mem);
 		if (ret)
 			goto undo;
 
-		ret = edgetpu_mailbox_alloc_queue(ext_mailbox->etdev, mailbox, attr.resp_queue_size,
-						  attr.sizeof_resp, MAILBOX_RESP_QUEUE,
-						  &desc->resp_queue_mem);
+		ret = edgetpu_mailbox_do_alloc_queue(etdev, mailbox, attr.resp_queue_size,
+						     attr.sizeof_resp, &desc->resp_queue_mem);
 		if (ret) {
-			edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox,
+			edgetpu_mailbox_free_queue(etdev, mailbox,
 						   &desc->cmd_queue_mem);
 			goto undo;
 		}
@@ -810,8 +856,8 @@ undo:
 	while (i--) {
 		desc = &ext_mailbox->descriptors[i];
 		mailbox = desc->mailbox;
-		edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox, &desc->cmd_queue_mem);
-		edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox, &desc->resp_queue_mem);
+		edgetpu_mailbox_free_queue(etdev, mailbox, &desc->cmd_queue_mem);
+		edgetpu_mailbox_free_queue(etdev, mailbox, &desc->resp_queue_mem);
 	}
 	return ret;
 }
@@ -821,12 +867,13 @@ static void edgetpu_mailbox_external_free_queue_batch(struct edgetpu_external_ma
 	u32 i;
 	struct edgetpu_mailbox *mailbox;
 	struct edgetpu_mailbox_descriptor *desc;
+	struct edgetpu_dev *etdev = ext_mailbox->etdev;
 
 	for (i = 0; i < ext_mailbox->count; i++) {
 		desc = &ext_mailbox->descriptors[i];
 		mailbox = desc->mailbox;
-		edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox, &desc->cmd_queue_mem);
-		edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox, &desc->resp_queue_mem);
+		edgetpu_mailbox_free_queue(etdev, mailbox, &desc->cmd_queue_mem);
+		edgetpu_mailbox_free_queue(etdev, mailbox, &desc->resp_queue_mem);
 	}
 }
 
@@ -840,6 +887,12 @@ static bool edgetpu_mailbox_external_check_range(struct edgetpu_mailbox_manager
 	return (start <= end) && (mgr->ext_index_from <= start && mgr->ext_index_to > end);
 }
 
+/*
+ * Allocates external mailboxes according to @ext_mailbox_req object and
+ * associate it with @group.
+ *
+ * Caller should hold @group->lock
+ */
 static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 					  struct edgetpu_external_mailbox_req *ext_mailbox_req)
 {
@@ -851,6 +904,12 @@ static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 	struct edgetpu_mailbox_attr attr;
 	unsigned long flags;
 
+	if (!edgetpu_device_group_is_finalized(group))
+		return -EINVAL;
+
+	if (group->ext_mailbox)
+		return -EEXIST;
+
 	if (!ext_mailbox_req)
 		return -EINVAL;
 
@@ -922,11 +981,7 @@ static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
 
 	for (i = 0; i < count; i++) {
 		mailbox = ext_mailbox->descriptors[i].mailbox;
-		edgetpu_mailbox_set_priority(mailbox, attr.priority);
-		EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, cmd_queue_tail_doorbell_enable,
-					      attr.cmdq_tail_doorbell);
 		mailbox->internal.group = edgetpu_device_group_get(group);
-		edgetpu_mailbox_enable(mailbox);
 	}
 	group->ext_mailbox = ext_mailbox;
 	return 0;
@@ -954,9 +1009,6 @@ static void edgetpu_mailbox_external_free(struct edgetpu_device_group *group)
 
 	mgr = ext_mailbox->etdev->mailbox_manager;
 
-	for (i = 0; i < ext_mailbox->count; i++)
-		edgetpu_mailbox_disable(ext_mailbox->descriptors[i].mailbox);
-
 	edgetpu_mailbox_external_free_queue_batch(ext_mailbox);
 
 	for (i = 0; i < ext_mailbox->count; i++)  {
@@ -973,11 +1025,8 @@ static void edgetpu_mailbox_external_free(struct edgetpu_device_group *group)
 static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 						 struct edgetpu_external_mailbox_req *req)
 {
-	int ret = 0, i;
-	struct edgetpu_external_mailbox *ext_mailbox = NULL;
+	int ret = 0;
 	struct edgetpu_device_group *group;
-	int vcid;
-	u32 mbox_map = 0;
 
 	mutex_lock(&client->group_lock);
 	if (!client->group || !edgetpu_device_group_is_leader(client->group, client)) {
@@ -987,51 +1036,27 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (!edgetpu_wakelock_lock(client->wakelock)) {
-		etdev_err(client->etdev, "Enabling external mailbox needs wakelock acquired\n");
-		edgetpu_wakelock_unlock(client->wakelock);
-		edgetpu_device_group_put(client->group);
-		return -EAGAIN;
-	}
-
-	mutex_lock(&group->lock);
-	if (!edgetpu_device_group_is_finalized(group)) {
-		ret = -EINVAL;
-		goto unlock;
-	}
-
-	if (group->ext_mailbox) {
-		ret = -EEXIST;
-		goto unlock;
-	}
-
-	ret = edgetpu_mailbox_external_alloc(group, req);
-	if (ret)
-		goto unlock;
-
-	ext_mailbox = group->ext_mailbox;
-	vcid = group->vcid;
-
-	for (i = 0; i < ext_mailbox->count; i++)
-		mbox_map |= BIT(ext_mailbox->descriptors[i].mailbox->mailbox_id);
-
-	ret = edgetpu_mailbox_activate_bulk(group->etdev, mbox_map, vcid, false);
-
-	if (ret) {
-		etdev_err(group->etdev, "Activate mailbox bulk failed: %d", ret);
-		/*
-		 * Deactivate only fails if f/w is unresponsive which will put group
-		 * in errored state or mailbox physically disabled before requesting
-		 * deactivate which will never be the case.
-		 */
-		edgetpu_mailbox_external_free(group);
+	if (edgetpu_pm_get_if_powered(group->etdev->pm)) {
+		mutex_lock(&group->lock);
+		ret = edgetpu_mailbox_external_alloc(group, req);
+		if (ret) {
+			mutex_unlock(&group->lock);
+			goto err;
+		}
+		edgetpu_mailbox_init_external_mailbox(group->ext_mailbox);
+		ret = edgetpu_mailbox_activate_external_mailbox(group);
+		mutex_unlock(&group->lock);
+		edgetpu_pm_put(group->etdev->pm);
+		goto out;
+	} else {
+		mutex_lock(&group->lock);
+		ret = edgetpu_mailbox_external_alloc(group, req);
+		mutex_unlock(&group->lock);
+		goto out;
 	}
-unlock:
-	mutex_unlock(&group->lock);
-	if (!ret)
-		edgetpu_wakelock_inc_event_locked(client->wakelock,
-						  EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
-	edgetpu_wakelock_unlock(client->wakelock);
+err:
+	edgetpu_pm_put(group->etdev->pm);
+out:
 	edgetpu_device_group_put(group);
 	return ret;
 }
@@ -1048,40 +1073,32 @@ static int edgetpu_mailbox_external_disable_free(struct edgetpu_client *client)
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (!edgetpu_wakelock_lock(client->wakelock)) {
-		etdev_err(client->etdev, "Disabling external mailbox needs wakelock acquired\n");
-		edgetpu_wakelock_unlock(client->wakelock);
-		edgetpu_device_group_put(client->group);
-		return -EAGAIN;
+	if (edgetpu_pm_get_if_powered(group->etdev->pm)) {
+		mutex_lock(&group->lock);
+		edgetpu_mailbox_external_disable_free_locked(group);
+		mutex_unlock(&group->lock);
+		edgetpu_pm_put(group->etdev->pm);
+	} else {
+		mutex_lock(&group->lock);
+		edgetpu_mailbox_external_free(group);
+		mutex_unlock(&group->lock);
 	}
-	mutex_lock(&group->lock);
-	edgetpu_mailbox_external_disable_free_locked(group);
-	mutex_unlock(&group->lock);
-	edgetpu_wakelock_dec_event_locked(client->wakelock, EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
-	edgetpu_wakelock_unlock(client->wakelock);
+
 	edgetpu_device_group_put(group);
 	return 0;
 }
 
 void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *group)
 {
-	u32 i, mbox_map = 0;
-	struct edgetpu_external_mailbox *ext_mailbox;
-
-	ext_mailbox = group->ext_mailbox;
-	if (!ext_mailbox)
-		return;
-
-	for (i = 0; i < ext_mailbox->count; i++)
-		mbox_map |= BIT(ext_mailbox->descriptors[i].mailbox->mailbox_id);
-
-	etdev_dbg(group->etdev, "Disabling mailboxes in map: %x\n", mbox_map);
-	edgetpu_mailbox_deactivate_bulk(group->etdev, mbox_map);
+	if (!group->dev_inaccessible) {
 	/*
 	 * Deactivate only fails if f/w is unresponsive which will put group
 	 * in errored state or mailbox physically disabled before requesting
 	 * deactivate which will never be the case.
 	 */
+		edgetpu_mailbox_deactivate_external_mailbox(group);
+		edgetpu_mailbox_disable_external_mailbox(group);
+	}
 	edgetpu_mailbox_external_free(group);
 }
 
@@ -1132,6 +1149,54 @@ static int edgetpu_mailbox_external_disable_by_id(struct edgetpu_client *client,
 	return ret;
 }
 
+int edgetpu_mailbox_activate_external_mailbox(struct edgetpu_device_group *group)
+{
+	struct edgetpu_external_mailbox *ext_mailbox = group->ext_mailbox;
+	uint vcid = group->vcid;
+	u32 mbox_map = 0, i;
+	int ret;
+
+	if (!ext_mailbox)
+		return -ENOENT;
+
+	for (i = 0; i < ext_mailbox->count; i++)
+		mbox_map |= BIT(ext_mailbox->descriptors[i].mailbox->mailbox_id);
+
+	ret = edgetpu_mailbox_activate_bulk(ext_mailbox->etdev, mbox_map, vcid, false);
+
+	if (ret)
+		etdev_err(group->etdev, "Activate mailbox bulk failed: %d", ret);
+	return ret;
+}
+
+void edgetpu_mailbox_disable_external_mailbox(struct edgetpu_device_group *group)
+{
+	u32 i;
+	struct edgetpu_external_mailbox *ext_mailbox = group->ext_mailbox;
+
+	if (!ext_mailbox)
+		return;
+
+	for (i = 0; i < ext_mailbox->count; i++)
+		edgetpu_mailbox_disable(ext_mailbox->descriptors[i].mailbox);
+
+}
+
+void edgetpu_mailbox_deactivate_external_mailbox(struct edgetpu_device_group *group)
+{
+	u32 i, mbox_map = 0;
+	struct edgetpu_external_mailbox *ext_mailbox = group->ext_mailbox;
+
+	if (!ext_mailbox)
+		return;
+
+	for (i = 0; i < ext_mailbox->count; i++)
+		mbox_map |= BIT(ext_mailbox->descriptors[i].mailbox->mailbox_id);
+
+	etdev_dbg(ext_mailbox->etdev, "Disabling mailboxes in map: %x\n", mbox_map);
+	edgetpu_mailbox_deactivate_bulk(ext_mailbox->etdev, mbox_map);
+}
+
 int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, int mailbox_id,
 			       struct edgetpu_external_mailbox_req *ext_mailbox_req)
 {
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index 11c9768..115341d 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -308,10 +308,10 @@ void edgetpu_mailbox_remove_vii(struct edgetpu_vii *vii);
 
 
 /*
- * Reset VII mailboxes CSRs to valid values, needed after the device is power
+ * Reset all mailboxes CSRs to valid values, needed after the device is power
  * gated.
  */
-void edgetpu_mailbox_reset_vii(struct edgetpu_mailbox_manager *mgr);
+void edgetpu_mailbox_reset_mailboxes(struct edgetpu_mailbox_manager *mgr);
 
 
 /* For VII and P2P mailboxes to allocate/free queue memory */
@@ -333,10 +333,10 @@ void edgetpu_mailbox_free_queue(struct edgetpu_dev *etdev,
 void edgetpu_mailbox_reinit_vii(struct edgetpu_device_group *group);
 
 /*
- * Re-configure VII mailbox queues which have an active client, re-using
+ * Re-configure VII and external mailbox queues which have an active client, re-using
  * existing buffers
  */
-void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev);
+void edgetpu_mailbox_restore_active_mailbox_queues(struct edgetpu_dev *etdev);
 
 /* utility functions for P2P */
 
@@ -392,10 +392,39 @@ void edgetpu_handshake_clear_fw_state(struct edgetpu_handshake *eh);
 /*
  * Disables and frees any external mailboxes allocated for @group.
  *
- * Caller must hold @group->lock.
+ * Caller must hold @group->lock and ensure if device is accessible.
  */
 void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *group);
 
+/*
+ * Re-programs the CSRs of queue addresses, context, priority etc. to @group's
+ * external mailbox.
+ *
+ * Caller holds @group->lock and ensures @group has mailbox attached.
+ */
+void edgetpu_mailbox_reinit_external_mailbox(struct edgetpu_device_group *group);
+
+/*
+ * Activates external mailboxes in @group's ext_mailbox.
+ *
+ * Caller ensures device is powered and must hold @group->lock.
+ */
+int edgetpu_mailbox_activate_external_mailbox(struct edgetpu_device_group *group);
+
+/*
+ * Deactivates external mailboxes in @group's ext_mailbox.
+ *
+ * Caller ensures device is powered and must hold @group->lock.
+ */
+void edgetpu_mailbox_deactivate_external_mailbox(struct edgetpu_device_group *group);
+
+/*
+ * Disables external mailboxes in @group's ext_mailbox.
+ *
+ * Caller ensures device is powered and must hold @group->lock.
+ */
+void edgetpu_mailbox_disable_external_mailbox(struct edgetpu_device_group *group);
+
 /* Utilities of circular queue operations */
 
 /*
diff --git a/drivers/edgetpu/edgetpu-mapping.h b/drivers/edgetpu/edgetpu-mapping.h
index 8f075aa..b022c87 100644
--- a/drivers/edgetpu/edgetpu-mapping.h
+++ b/drivers/edgetpu/edgetpu-mapping.h
@@ -165,11 +165,11 @@ static inline int __dma_dir_to_iommu_prot(enum dma_data_direction dir)
 static inline int mmu_flag_to_iommu_prot(u32 mmu_flags, struct device *dev,
 					 enum dma_data_direction dir)
 {
-	int prot = 0; /* hardcode to non-dma-coherent for prior kernels */
+	int prot = 0;
 
 	if (mmu_flags & EDGETPU_MMU_COHERENT) {
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 10, 0)
-		prot = dev_is_dma_coherent(dev) ? IOMMU_CACHE : 0;
+#ifdef EDGETPU_IS_DMA_COHERENT
+		prot = IOMMU_CACHE;
 #endif
 	}
 	prot |= __dma_dir_to_iommu_prot(dir);
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index e6bfa1d..ae075d0 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -55,7 +55,7 @@ static int edgetpu_pm_get_locked(struct edgetpu_pm *etpm)
 	if (!power_up_count) {
 		ret = etpm->p->handlers->power_up(etpm);
 		if (!ret)
-			edgetpu_mailbox_restore_active_vii_queues(etpm->etdev);
+			edgetpu_mailbox_restore_active_mailbox_queues(etpm->etdev);
 	}
 	if (ret)
 		etpm->p->power_up_count--;
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index a21bed2..52a43ab 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -449,8 +449,8 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 		edgetpu_kci_reinit(etdev->kci);
 	}
 	if (etdev->mailbox_manager) {
-		etdev_dbg(etdev, "Resetting VII mailboxes\n");
-		edgetpu_mailbox_reset_vii(etdev->mailbox_manager);
+		etdev_dbg(etdev, "Resetting (VII/external) mailboxes\n");
+		edgetpu_mailbox_reset_mailboxes(etdev->mailbox_manager);
 	}
 
 	if (!etdev->firmware)
@@ -741,7 +741,7 @@ static void mobile_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 	mutex_unlock(&platform_pwr->scenario_lock);
 }
 
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val)
+void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
 {
 	etdev_dbg(etdev, "%s: bts request - val = %u\n", __func__, bts_val);
 
diff --git a/drivers/edgetpu/mobile-pm.h b/drivers/edgetpu/mobile-pm.h
index 0cb4b77..4299b49 100644
--- a/drivers/edgetpu/mobile-pm.h
+++ b/drivers/edgetpu/mobile-pm.h
@@ -47,6 +47,7 @@ static inline int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
 enum mobile_reverse_kci_code {
 	RKCI_CODE_PM_QOS = RKCI_CHIP_CODE_FIRST + 1,
 	RKCI_CODE_BTS = RKCI_CHIP_CODE_FIRST + 2,
+	RKCI_CODE_PM_QOS_BTS = RKCI_CHIP_CODE_FIRST + 3,
 };
 
 #define MAX_VOLTAGE_VAL 1250000
@@ -84,6 +85,6 @@ void mobile_pm_destroy(struct edgetpu_dev *etdev);
 void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
 
 /* Set BTS value for the edgetpu device. */
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val);
+void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val);
 
 #endif /* __MOBILE_PM_H__ */
diff --git a/drivers/edgetpu/mobile-thermal.c b/drivers/edgetpu/mobile-thermal.c
index c1b1f69..09705a8 100644
--- a/drivers/edgetpu/mobile-thermal.c
+++ b/drivers/edgetpu/mobile-thermal.c
@@ -16,7 +16,6 @@
 #include <linux/slab.h>
 #include <linux/thermal.h>
 #include <linux/version.h>
-#include <soc/google/gs101_tmu.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
@@ -61,6 +60,26 @@ static int edgetpu_get_max_state(struct thermal_cooling_device *cdev, unsigned l
 	return 0;
 }
 
+static int edgetpu_set_thermal_policy(struct edgetpu_dev *etdev, unsigned long pwr_state)
+{
+	int ret;
+
+	if (!edgetpu_pm_trylock(etdev->pm))
+		return -EAGAIN;
+
+	if (edgetpu_is_powered(etdev))
+		edgetpu_kci_block_bus_speed_control(etdev, true);
+
+	ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, pwr_state);
+
+	if (edgetpu_is_powered(etdev))
+		edgetpu_kci_block_bus_speed_control(etdev, false);
+
+	edgetpu_pm_unlock(etdev->pm);
+
+	return ret;
+}
+
 /*
  * Set cooling state.
  */
@@ -93,9 +112,9 @@ static int edgetpu_set_cur_state(struct thermal_cooling_device *cdev, unsigned l
 	if (pwr_state < TPU_ACTIVE_UUD) {
 		dev_warn_ratelimited(dev,
 				     "Setting lowest DVFS state, waiting for FW to shutdown TPU");
-		ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, TPU_ACTIVE_UUD);
+		ret = edgetpu_set_thermal_policy(cooling->etdev, TPU_ACTIVE_UUD);
 	} else {
-		ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, pwr_state);
+		ret = edgetpu_set_thermal_policy(cooling->etdev, pwr_state);
 	}
 
 	if (ret) {