[Copybara Auto Merge] Merge branch 'whitechapel' into android-gs-pixel-5.10

  edgetpu: abrolhos return actual error for throttling kci
  edgetpu: handle job lockup notification from firmware
  edgetpu: add API to send fatal error notifications to a specific group
  edgetpu: add fatal error event for firmware-detected job timeout
  edgetpu: Improve the check in edgetpu_thermal_* functions
  edgetpu: Modify the log function due to thermal suspended
  edgetpu: fail wakelock acquiring if suspended
  edgetpu: ignore offset arg in edgetpu_map_dmabuf
  edgetpu: abrolhos: hook exynos acpm functions
  edgetpu: abrolhos remove memset log/trace buffers
  edgetpu: update mobile SoC DT bindings
  edgetpu: mapping remove WARN when release not set
  edgetpu: remove reference to future chipset name in common source
  edgetpu: allocate/free external mailboxes
  edgetpu: add support for external mailboxes
  edgetpu: KCI allocate data with iremap

GitOrigin-RevId: 7374471f6dd6c760248d0cf03c6bb3dacc0fdf1d
Change-Id: I48ecd34e0c3378dcb9f3192eeda969e8d7923c21

20 files changed, 539 insertions, 118 deletions

diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index fe3da83..1277199 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -207,7 +207,7 @@ int edgetpu_chip_acquire_ext_mailbox(struct edgetpu_client *client,
 		return -EBUSY;
 	}
 	apdev->secure_client = client;
-	ret = edgetpu_mailbox_enable_ext(client, ABROLHOS_TZ_MAILBOX_ID);
+	ret = edgetpu_mailbox_enable_ext(client, ABROLHOS_TZ_MAILBOX_ID, NULL);
 	mutex_unlock(&apdev->tz_mailbox_lock);
 	return ret;
 }
diff --git a/drivers/edgetpu/abrolhos-platform.c b/drivers/edgetpu/abrolhos-platform.c
index caf0910..753e18d 100644
--- a/drivers/edgetpu/abrolhos-platform.c
+++ b/drivers/edgetpu/abrolhos-platform.c
@@ -32,9 +32,12 @@
 #include "mobile-firmware.h"
 
 static const struct of_device_id edgetpu_of_match[] = {
+	/* TODO(b/190677977): remove  */
 	{ .compatible = "google,darwinn", },
+	{ .compatible = "google,edgetpu-gs101", },
 	{ /* end of list */ },
 };
+
 MODULE_DEVICE_TABLE(of, edgetpu_of_match);
 
 static void sscd_release(struct device *dev)
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 2930917..9a3db2f 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -22,7 +22,6 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 #include "edgetpu-pm.h"
-#include "edgetpu-telemetry.h"
 
 #include "edgetpu-pm.c"
 
@@ -445,12 +444,6 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 	if (ret)
 		return ret;
 
-	/* Clear out log / trace buffers */
-	memset(abpdev->log_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
-#if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-	memset(abpdev->trace_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
-#endif
-
 	edgetpu_chip_init(etdev);
 
 	if (etdev->kci) {
diff --git a/drivers/edgetpu/abrolhos-pm.h b/drivers/edgetpu/abrolhos-pm.h
index 0f7bdf0..48b65be 100644
--- a/drivers/edgetpu/abrolhos-pm.h
+++ b/drivers/edgetpu/abrolhos-pm.h
@@ -11,7 +11,7 @@
 #include "edgetpu-internal.h"
 
 /* Can't build out of tree with acpm_dvfs unless kernel supports ACPM */
-#if IS_ENABLED(CONFIG_ACPM_DVFS)
+#if IS_ENABLED(CONFIG_ACPM_DVFS) || IS_ENABLED(CONFIG_EDGETPU_TEST)
 
 #include <linux/acpm_dvfs.h>
 
@@ -37,7 +37,7 @@ static inline int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
 {
 	return 0;
 }
-#endif /* IS_ENABLED(CONFIG_ACPM_DVFS) */
+#endif /* IS_ENABLED(CONFIG_ACPM_DVFS) || IS_ENABLED(CONFIG_EDGETPU_TEST) */
 
 /*
  * TPU Power States:
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index 572518b..d8123b6 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -77,17 +77,15 @@ static int edgetpu_set_cur_state(struct thermal_cooling_device *cdev,
 			dev_err(dev, "error setting tpu policy: %d\n", ret);
 			goto out;
 		}
+		if (state_original == 0)
+			cooling->thermal_suspended = true;
+		else if (cooling->cooling_state == 0)
+			cooling->thermal_suspended = false;
 		cooling->cooling_state = state_original;
 		ret = edgetpu_kci_notify_throttling(etdev, pwr_state);
-		if (ret) {
-			/*
-			 * TODO(b/185596886) : After FW adds a handler for this KCI, return the
-			 * correct value of ret and change the debug message to an error message.
-			 */
-			etdev_dbg(etdev, "Failed to notify FW about state %lu, error:%d",
+		if (ret)
+			etdev_err(etdev, "Failed to notify FW about power state %lu, error:%d",
 				  pwr_state, ret);
-			ret = 0;
-		}
 	} else {
 		ret = -EALREADY;
 	}
diff --git a/drivers/edgetpu/abrolhos/config-mailbox.h b/drivers/edgetpu/abrolhos/config-mailbox.h
index 4bec2d5..efbd169 100644
--- a/drivers/edgetpu/abrolhos/config-mailbox.h
+++ b/drivers/edgetpu/abrolhos/config-mailbox.h
@@ -13,6 +13,7 @@
 #define EDGETPU_NUM_MAILBOXES 8
 #define EDGETPU_NUM_VII_MAILBOXES (EDGETPU_NUM_MAILBOXES - 1)
 #define EDGETPU_NUM_P2P_MAILBOXES 0
+#define EDGETPU_NUM_EXT_MAILBOXES 0
 
 #define ABROLHOS_TZ_MAILBOX_ID			8
 
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 86bf7d8..7701a71 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -323,6 +323,7 @@ static struct edgetpu_mailbox_manager_desc mailbox_manager_desc = {
 	.num_mailbox = EDGETPU_NUM_MAILBOXES,
 	.num_vii_mailbox = EDGETPU_NUM_VII_MAILBOXES,
 	.num_p2p_mailbox = EDGETPU_NUM_P2P_MAILBOXES,
+	.num_ext_mailbox = EDGETPU_NUM_EXT_MAILBOXES,
 	.get_context_csr_base = edgetpu_mailbox_get_context_csr_base,
 	.get_cmd_queue_csr_base = edgetpu_mailbox_get_cmd_queue_csr_base,
 	.get_resp_queue_csr_base = edgetpu_mailbox_get_resp_queue_csr_base,
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 97f3254..53c8ca2 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -475,6 +475,7 @@ static void edgetpu_device_group_release(struct edgetpu_device_group *group)
 #ifdef EDGETPU_HAS_P2P_MAILBOX
 		edgetpu_p2p_mailbox_release(group);
 #endif
+		edgetpu_mailbox_external_disable_free_locked(group);
 		edgetpu_mailbox_remove_vii(&group->vii);
 		group_release_members(group);
 	}
@@ -1676,6 +1677,27 @@ out:
 }
 
 /*
+ * Set @group status as errored, set the error mask, and notify the runtime of
+ * the fatal error event on the group.
+ */
+void edgetpu_group_fatal_error_notify(struct edgetpu_device_group *group,
+				      uint error_mask)
+{
+	etdev_dbg(group->etdev, "notify group %u error 0x%x",
+		  group->workload_id, error_mask);
+	mutex_lock(&group->lock);
+	/*
+	 * Only finalized groups may have handshake with the FW, mark
+	 * them as errored.
+	 */
+	if (edgetpu_device_group_is_finalized(group))
+		group->status = EDGETPU_DEVICE_GROUP_ERRORED;
+	group->fatal_errors |= error_mask;
+	mutex_unlock(&group->lock);
+	edgetpu_group_notify(group, EDGETPU_EVENT_FATAL_ERROR);
+}
+
+/*
  * For each group active on @etdev: set the group status as errored, set the
  * error mask, and notify the runtime of the fatal error event.
  */
@@ -1708,19 +1730,8 @@ void edgetpu_fatal_error_notify(struct edgetpu_dev *etdev, uint error_mask)
 	}
 	mutex_unlock(&etdev->groups_lock);
 	for (i = 0; i < num_groups; i++) {
-		group = groups[i];
-		mutex_lock(&group->lock);
-		/*
-		 * Only finalized groups may have handshake with the FW, mark
-		 * them as errored.
-		 */
-
-		if (edgetpu_device_group_is_finalized(group))
-			group->status = EDGETPU_DEVICE_GROUP_ERRORED;
-		group->fatal_errors |= error_mask;
-		mutex_unlock(&group->lock);
-		edgetpu_group_notify(group, EDGETPU_EVENT_FATAL_ERROR);
-		edgetpu_device_group_put(group);
+		edgetpu_group_fatal_error_notify(groups[i], error_mask);
+		edgetpu_device_group_put(groups[i]);
 	}
 	kfree(groups);
 }
@@ -1785,3 +1796,41 @@ int edgetpu_group_attach_and_open_mailbox(struct edgetpu_device_group *group)
 	mutex_unlock(&group->lock);
 	return ret;
 }
+
+/*
+ * Return the group with id @vcid for device @etdev, with a reference held
+ * on the group (must call edgetpu_device_group_put when done), or NULL if
+ * no group with that VCID is found.
+ */
+static struct edgetpu_device_group *get_group_by_vcid(
+	struct edgetpu_dev *etdev, u16 vcid)
+{
+	struct edgetpu_device_group *group = NULL;
+	struct edgetpu_device_group *tgroup;
+	struct edgetpu_list_group *g;
+
+	mutex_lock(&etdev->groups_lock);
+	etdev_for_each_group(etdev, g, tgroup) {
+		if (tgroup->vcid == vcid) {
+			group = edgetpu_device_group_get(tgroup);
+			break;
+		}
+	}
+	mutex_unlock(&etdev->groups_lock);
+	return group;
+}
+
+void edgetpu_handle_job_lockup(struct edgetpu_dev *etdev, u16 vcid)
+{
+	struct edgetpu_device_group *group;
+
+	etdev_err(etdev, "firmware-detected job lockup on VCID %u",
+		  vcid);
+	group = get_group_by_vcid(etdev, vcid);
+	if (!group) {
+		etdev_warn(etdev, "VCID %u group not found", vcid);
+		return;
+	}
+	edgetpu_group_fatal_error_notify(group, EDGETPU_ERROR_RUNTIME_TIMEOUT);
+	edgetpu_device_group_put(group);
+}
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 5e05799..7b20dd0 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -120,6 +120,11 @@ struct edgetpu_device_group {
 	struct edgetpu_iommu_domain *etdomain;
 	/* matrix of P2P mailboxes */
 	struct edgetpu_p2p_mailbox **p2p_mailbox_matrix;
+	/*
+	 * External mailboxes associated with this group, only valid if
+	 * external mailbox allocated and enabled.
+	 */
+	struct edgetpu_external_mailbox *ext_mailbox;
 
 	/* Mask of errors set for this group. */
 	uint fatal_errors;
@@ -378,6 +383,9 @@ bool edgetpu_in_any_group(struct edgetpu_dev *etdev);
  */
 bool edgetpu_set_group_join_lockout(struct edgetpu_dev *etdev, bool lockout);
 
+/* Notify @group about a fatal error for that group. */
+void edgetpu_group_fatal_error_notify(struct edgetpu_device_group *group,
+				      uint error_mask);
 /* Notify all device groups of @etdev about a failure on the die */
 void edgetpu_fatal_error_notify(struct edgetpu_dev *etdev, uint error_mask);
 
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 3d1c190..5259650 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -36,10 +36,8 @@ struct dmabuf_map_entry {
 	/* SG table returned by dma_buf_map_attachment() */
 	struct sg_table *sgt;
 	/*
-	 * The SG table that shrunk from @sgt with region [offset, offset+size],
-	 * where @offset and @size are the arguments in edgetpu_dmabuf_map.
-	 * If @offset for mapping is zero and @size equals the total length of
-	 * @sgt, this table is a duplicate of @sgt.
+	 * The SG table that shrunk and condensed from @sgt with region [0, size), where @size is
+	 * the size field in edgetpu_dmabuf_map which owns this entry.
 	 */
 	struct sg_table shrunk_sgt;
 };
@@ -50,7 +48,6 @@ struct dmabuf_map_entry {
  */
 struct edgetpu_dmabuf_map {
 	struct edgetpu_mapping map;
-	u64 offset;
 	u64 size; /* size of this mapping in bytes */
 	u32 mmu_flags;
 	/*
@@ -285,7 +282,6 @@ static void dmabuf_map_callback_release(struct edgetpu_mapping *map)
 	uint i;
 
 	if (tpu_addr) {
-		tpu_addr -= dmap->offset;
 		if (IS_MIRRORED(map->flags)) {
 			group_unmap_dmabuf(group, dmap, tpu_addr);
 		} else {
@@ -636,7 +632,6 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	int ret = -EINVAL;
 	struct dma_buf *dmabuf;
 	edgetpu_map_flag_t flags = arg->flags;
-	const u64 offset = arg->offset;
 	u64 size;
 	const enum dma_data_direction dir =
 		edgetpu_host_dma_dir(flags & EDGETPU_MAP_DIR_MASK);
@@ -645,30 +640,16 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	tpu_addr_t tpu_addr;
 	uint i;
 
-	/* invalid DMA direction or offset is not page-aligned */
-	if (!valid_dma_direction(dir) || offset_in_page(offset)) {
-		etdev_dbg(group->etdev,
-			  "%s: valid=%d offset_in_page=%lu offset=0x%llx\n",
-			  __func__, valid_dma_direction(dir),
-			  offset_in_page(offset), offset);
+	if (!valid_dma_direction(dir)) {
+		etdev_dbg(group->etdev, "%s: invalid direction %d\n", __func__, dir);
 		return -EINVAL;
 	}
-	/* TODO(b/189278468): entirely ignore @offset */
-	if (offset != 0)
-		etdev_warn_ratelimited(group->etdev,
-				       "Non-zero offset for dmabuf mapping is deprecated");
 	dmabuf = dma_buf_get(arg->dmabuf_fd);
 	if (IS_ERR(dmabuf)) {
 		etdev_dbg(group->etdev, "%s: dma_buf_get returns %ld\n",
 			  __func__, PTR_ERR(dmabuf));
 		return PTR_ERR(dmabuf);
 	}
-	if (offset >= dmabuf->size) {
-		etdev_dbg(group->etdev,
-			  "%s: offset=0x%llx > dmabuf size=%zx\n",
-			  __func__, offset, dmabuf->size);
-		goto err_put;
-	}
 
 	mutex_lock(&group->lock);
 	if (!edgetpu_device_group_is_finalized(group)) {
@@ -687,7 +668,6 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 
 	get_dma_buf(dmabuf);
 	dmap->dmabufs[0] = dmabuf;
-	dmap->offset = offset;
 	dmap->size = size = dmabuf->size;
 	if (IS_MIRRORED(flags)) {
 		for (i = 0; i < group->n_clients; i++) {
@@ -734,7 +714,7 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 		}
 		dmap->map.die_index = arg->die_index;
 	}
-	dmap->map.device_address = tpu_addr + offset;
+	dmap->map.device_address = tpu_addr;
 	ret = edgetpu_mapping_add(&group->dmabuf_mappings, &dmap->map);
 	if (ret) {
 		etdev_dbg(group->etdev, "%s: edgetpu_mapping_add returns %d\n",
@@ -751,7 +731,6 @@ err_release_map:
 	dmabuf_map_callback_release(&dmap->map);
 err_unlock_group:
 	mutex_unlock(&group->lock);
-err_put:
 	dma_buf_put(dmabuf);
 
 	return ret;
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 311ea86..6fbd642 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -579,6 +579,7 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 {
 	int count;
 	int ret;
+	struct edgetpu_thermal *thermal = client->etdev->thermal;
 
 	edgetpu_wakelock_lock(client->wakelock);
 	/* when NO_WAKELOCK: count should be 1 so here is a no-op */
@@ -588,7 +589,18 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 		return count;
 	}
 	if (!count) {
-		ret = edgetpu_pm_get(client->etdev->pm);
+		edgetpu_thermal_lock(thermal);
+		if (edgetpu_thermal_is_suspended(thermal)) {
+			/* TPU is thermal suspended, so fail acquiring wakelock */
+			ret = -EAGAIN;
+			etdev_warn_ratelimited(client->etdev,
+					       "wakelock acquire rejected due to thermal suspend");
+			edgetpu_thermal_unlock(thermal);
+			goto error_release;
+		} else {
+			ret = edgetpu_pm_get(client->etdev->pm);
+			edgetpu_thermal_unlock(thermal);
+		}
 		if (ret) {
 			etdev_warn(client->etdev, "%s: pm_get failed (%d)",
 				   __func__, ret);
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index f5c6853..7c4966e 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -339,6 +339,9 @@ static inline bool edgetpu_is_external_wrapper_class_file(struct file *file)
 void edgetpu_handle_firmware_crash(struct edgetpu_dev *etdev,
 				   enum edgetpu_fw_crash_type crash_type);
 
+/* Handle notification of job lockup from firmware */
+void edgetpu_handle_job_lockup(struct edgetpu_dev *etdev, u16 vcid);
+
 /* Bus (Platform/PCI) <-> Core API */
 
 int __init edgetpu_init(void);
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 0adf9a5..73a47cc 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -9,7 +9,6 @@
 #include <linux/bits.h>
 #include <linux/circ_buf.h>
 #include <linux/device.h>
-#include <linux/dma-mapping.h> /* dmam_alloc_coherent */
 #include <linux/errno.h>
 #include <linux/seq_file.h>
 #include <linux/slab.h>
@@ -94,6 +93,9 @@ edgetpu_reverse_kci_consume_response(struct edgetpu_dev *etdev,
 		edgetpu_handle_firmware_crash(
 		      etdev, (enum edgetpu_fw_crash_type)resp->retval);
 		break;
+	case RKCI_JOB_LOCKUP:
+		edgetpu_handle_job_lockup(etdev, resp->retval);
+		break;
 	default:
 		etdev_warn(etdev, "%s: Unrecognized KCI request: 0x%x\n",
 			   __func__, resp->code);
@@ -696,33 +698,23 @@ static int edgetpu_kci_send_cmd_with_data(struct edgetpu_kci *kci,
 					  size_t size)
 {
 	struct edgetpu_dev *etdev = kci->mailbox->etdev;
-	dma_addr_t dma_addr;
-	tpu_addr_t tpu_addr;
+	struct edgetpu_coherent_mem mem;
 	int ret;
-	void *ptr = dma_alloc_coherent(etdev->dev, size, &dma_addr, GFP_KERNEL);
-	const u32 flags = EDGETPU_MMU_DIE | EDGETPU_MMU_32 | EDGETPU_MMU_HOST;
 
-	if (!ptr)
-		return -ENOMEM;
-	memcpy(ptr, data, size);
-
-	tpu_addr = edgetpu_mmu_tpu_map(etdev, dma_addr, size, DMA_TO_DEVICE, EDGETPU_CONTEXT_KCI,
-				       flags);
-	if (!tpu_addr) {
-		etdev_err(etdev, "%s: failed to map to TPU", __func__);
-		dma_free_coherent(etdev->dev, size, ptr, dma_addr);
-		return -ENOSPC;
-	}
-	etdev_dbg(etdev, "%s: map kva=%pK iova=0x%llx dma=%pad", __func__, ptr, tpu_addr,
-		  &dma_addr);
+	ret = edgetpu_iremap_alloc(etdev, size, &mem, EDGETPU_CONTEXT_KCI);
+	if (ret)
+		return ret;
+	memcpy(mem.vaddr, data, size);
 
-	cmd->dma.address = tpu_addr;
+	etdev_dbg(etdev, "%s: map kva=%pK iova=0x%llx dma=%pad", __func__, mem.vaddr, mem.tpu_addr,
+		  &mem.dma_addr);
+
+	cmd->dma.address = mem.tpu_addr;
 	cmd->dma.size = size;
 	ret = edgetpu_kci_send_cmd(kci, cmd);
-	edgetpu_mmu_tpu_unmap(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	dma_free_coherent(etdev->dev, size, ptr, dma_addr);
-	etdev_dbg(etdev, "%s: unmap kva=%pK iova=0x%llx dma=%pad", __func__, ptr, tpu_addr,
-		  &dma_addr);
+	edgetpu_iremap_free(etdev, &mem, EDGETPU_CONTEXT_KCI);
+	etdev_dbg(etdev, "%s: unmap kva=%pK iova=0x%llx dma=%pad", __func__, mem.vaddr,
+		  mem.tpu_addr, &mem.dma_addr);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index 97e4079..deb258d 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -122,6 +122,7 @@ enum edgetpu_reverse_kci_code {
 	RKCI_CHIP_CODE_LAST = 0x7FFF,
 	RKCI_GENERIC_CODE_FIRST = 0x8000,
 	RKCI_FIRMWARE_CRASH = RKCI_GENERIC_CODE_FIRST + 0,
+	RKCI_JOB_LOCKUP = RKCI_GENERIC_CODE_FIRST + 1,
 	RKCI_GENERIC_CODE_LAST = 0xFFFF,
 };
 
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index ba33eab..cf996f7 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -80,6 +80,24 @@ edgetpu_mailbox_create_locked(struct edgetpu_mailbox_manager *mgr, uint index)
 	return mailbox;
 }
 
+/* Caller must hold @mgr->mailboxes_lock. */
+static int edgetpu_mailbox_remove_locked(struct edgetpu_mailbox_manager *mgr,
+					 struct edgetpu_mailbox *mailbox)
+{
+	/* simple security checks */
+	if (mailbox->mailbox_id >= mgr->num_mailbox ||
+	    mgr->mailboxes[mailbox->mailbox_id] != mailbox) {
+		return -EINVAL;
+	}
+
+	mgr->mailboxes[mailbox->mailbox_id] = NULL;
+	/* KCI mailbox is a special case */
+	if (mailbox->mailbox_id == KERNEL_MAILBOX_INDEX)
+		edgetpu_kci_release(mgr->etdev, mailbox->internal.kci);
+	kfree(mailbox);
+	return 0;
+}
+
 /*
  * Disables the @index-th mailbox via setting CSR. Doesn't need
  * @mgr->mailboxes[index] be allocated.
@@ -342,23 +360,13 @@ out:
 int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr, struct edgetpu_mailbox *mailbox)
 {
 	unsigned long flags;
+	int ret;
 
 	write_lock_irqsave(&mgr->mailboxes_lock, flags);
-	/* simple security checks */
-	if (mailbox->mailbox_id >= mgr->num_mailbox ||
-	    mgr->mailboxes[mailbox->mailbox_id] != mailbox) {
-		write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
-		return -EINVAL;
-	}
-
-	mgr->mailboxes[mailbox->mailbox_id] = NULL;
-	/* KCI mailbox is a special case */
-	if (mailbox->mailbox_id == KERNEL_MAILBOX_INDEX)
-		edgetpu_kci_release(mgr->etdev, mailbox->internal.kci);
+	ret = edgetpu_mailbox_remove_locked(mgr, mailbox);
 	write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
-	kfree(mailbox);
 
-	return 0;
+	return ret;
 }
 
 /*
@@ -548,6 +556,7 @@ edgetpu_mailbox_create_mgr(struct edgetpu_dev *etdev,
 	total += 1; /* KCI mailbox */
 	total += desc->num_vii_mailbox;
 	total += desc->num_p2p_mailbox;
+	total += desc->num_ext_mailbox;
 	if (total > desc->num_mailbox)
 		return ERR_PTR(-EINVAL);
 	mgr = devm_kzalloc(etdev->dev, sizeof(*mgr), GFP_KERNEL);
@@ -561,6 +570,8 @@ edgetpu_mailbox_create_mgr(struct edgetpu_dev *etdev,
 	mgr->vii_index_to = mgr->vii_index_from + desc->num_vii_mailbox;
 	mgr->p2p_index_from = mgr->vii_index_to;
 	mgr->p2p_index_to = mgr->p2p_index_from + desc->num_p2p_mailbox;
+	mgr->ext_index_from = mgr->p2p_index_to;
+	mgr->ext_index_to = mgr->ext_index_from + desc->num_ext_mailbox;
 
 	mgr->get_context_csr_base = desc->get_context_csr_base;
 	mgr->get_cmd_queue_csr_base = desc->get_cmd_queue_csr_base;
@@ -769,52 +780,340 @@ void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev)
 	kfree(groups);
 }
 
-int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, u32 mailbox_id)
+static int edgetpu_mailbox_external_alloc_queue_batch(struct edgetpu_external_mailbox *ext_mailbox)
+{
+	int ret, i;
+	struct edgetpu_mailbox *mailbox;
+	struct edgetpu_mailbox_attr attr;
+	struct edgetpu_mailbox_descriptor *desc;
+
+	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);
+		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);
+		if (ret) {
+			edgetpu_mailbox_free_queue(ext_mailbox->etdev, mailbox,
+						   &desc->cmd_queue_mem);
+			goto undo;
+		}
+	}
+	return 0;
+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);
+	}
+	return ret;
+}
+
+static void edgetpu_mailbox_external_free_queue_batch(struct edgetpu_external_mailbox *ext_mailbox)
+{
+	u32 i;
+	struct edgetpu_mailbox *mailbox;
+	struct edgetpu_mailbox_descriptor *desc;
+
+	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);
+	}
+}
+
+/*
+ * Checks if the indexes given for external mailboxes are in range of mailbox
+ * manager(@mgr) managing the external mailboxes.
+ */
+static bool edgetpu_mailbox_external_check_range(struct edgetpu_mailbox_manager *mgr,
+						 const int start, const int end)
+{
+	return (start <= end) && (mgr->ext_index_from <= start && mgr->ext_index_to > end);
+}
+
+static int edgetpu_mailbox_external_alloc(struct edgetpu_device_group *group,
+					  struct edgetpu_external_mailbox_req *ext_mailbox_req)
+{
+	u32 i, j = 0;
+	struct edgetpu_mailbox_manager *mgr = group->etdev->mailbox_manager;
+	struct edgetpu_mailbox *mailbox;
+	int ret = 0, c = 0, count;
+	struct edgetpu_external_mailbox *ext_mailbox;
+	struct edgetpu_mailbox_attr attr;
+	unsigned long flags;
+
+	if (!ext_mailbox_req)
+		return -EINVAL;
+
+	count = ext_mailbox_req->count;
+	attr = ext_mailbox_req->attr;
+
+	if (!edgetpu_mailbox_external_check_range(mgr, ext_mailbox_req->start,
+						  ext_mailbox_req->end))
+		return -ERANGE;
+
+	ext_mailbox = kzalloc(sizeof(*ext_mailbox), GFP_KERNEL);
+	if (!ext_mailbox)
+		return -ENOMEM;
+
+	ext_mailbox->descriptors =
+		kcalloc(count, sizeof(struct edgetpu_mailbox_descriptor), GFP_KERNEL);
+	if (!ext_mailbox->descriptors) {
+		kfree(ext_mailbox);
+		return -ENOMEM;
+	}
+
+	ext_mailbox->attr = attr;
+	ext_mailbox->count = count;
+	ext_mailbox->etdev = group->etdev;
+
+	write_lock_irqsave(&mgr->mailboxes_lock, flags);
+	for (i = ext_mailbox_req->start; i <= ext_mailbox_req->end; i++) {
+		if (!mgr->mailboxes[i])
+			c++;
+	}
+	if (c < count) {
+		ret = -EBUSY;
+		goto unlock;
+	}
+
+	for (i = ext_mailbox_req->start; i <= ext_mailbox_req->end && j < count; i++) {
+		if (!mgr->mailboxes[i]) {
+			mailbox = edgetpu_mailbox_create_locked(mgr, i);
+			if (!IS_ERR(mailbox)) {
+				mgr->mailboxes[i] = mailbox;
+				ext_mailbox->descriptors[j++].mailbox = mailbox;
+			} else {
+				goto release;
+			}
+		}
+	}
+
+	ret = edgetpu_mailbox_external_alloc_queue_batch(ext_mailbox);
+	if (ret)
+		goto release;
+	write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
+
+	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;
+release:
+	while (j--)
+		edgetpu_mailbox_remove_locked(mgr, ext_mailbox->descriptors[j].mailbox);
+unlock:
+	write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
+	kfree(ext_mailbox->descriptors);
+	kfree(ext_mailbox);
+	return ret;
+}
+
+/* Caller must hold @group->lock. */
+static void edgetpu_mailbox_external_free(struct edgetpu_device_group *group)
+{
+	struct edgetpu_mailbox_manager *mgr;
+	struct edgetpu_mailbox *mailbox;
+	struct edgetpu_external_mailbox *ext_mailbox;
+	u32 i;
+
+	ext_mailbox = group->ext_mailbox;
+	if (!ext_mailbox)
+		return;
+
+	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++)  {
+		mailbox = ext_mailbox->descriptors[i].mailbox;
+		edgetpu_device_group_put(mailbox->internal.group);
+		edgetpu_mailbox_remove(mgr, mailbox);
+	}
+
+	kfree(ext_mailbox->descriptors);
+	kfree(ext_mailbox);
+	group->ext_mailbox = NULL;
+}
+
+static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
+						 struct edgetpu_external_mailbox_req *req)
+{
+	int ret = 0, i, id;
+	struct edgetpu_external_mailbox *ext_mailbox = NULL;
+	struct edgetpu_device_group *group;
+	int vcid;
+
+	mutex_lock(&client->group_lock);
+	if (!client->group || !edgetpu_device_group_is_leader(client->group, client)) {
+		mutex_unlock(&client->group_lock);
+		return -EINVAL;
+	}
+	group = edgetpu_device_group_get(client->group);
+	mutex_unlock(&client->group_lock);
+
+	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++) {
+		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, false);
+		if (ret) {
+			etdev_err(group->etdev, "Activate mailbox %d failed: %d", id, ret);
+			break;
+		}
+	}
+
+	if (ret) {
+		while (i--) {
+			id = ext_mailbox->descriptors[i].mailbox->mailbox_id;
+			if (edgetpu_mailbox_deactivate(group->etdev, id))
+				etdev_err(group->etdev, "Deactivate mailbox %d failed", id);
+		}
+		/*
+		 * 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);
+	}
+unlock:
+	mutex_unlock(&group->lock);
+	edgetpu_device_group_put(group);
+	return ret;
+}
+
+static int edgetpu_mailbox_external_disable_free(struct edgetpu_client *client)
+{
+	struct edgetpu_device_group *group;
+
+	mutex_lock(&client->group_lock);
+	if (!client->group || !edgetpu_device_group_is_leader(client->group, client)) {
+		mutex_unlock(&client->group_lock);
+		return -EINVAL;
+	}
+	group = edgetpu_device_group_get(client->group);
+	mutex_unlock(&client->group_lock);
+
+	mutex_lock(&group->lock);
+	edgetpu_mailbox_external_disable_free_locked(group);
+	mutex_unlock(&group->lock);
+	edgetpu_device_group_put(group);
+	return 0;
+}
+
+void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *group)
+{
+	u32 i, id;
+	struct edgetpu_external_mailbox *ext_mailbox;
+
+	ext_mailbox = group->ext_mailbox;
+	if (!ext_mailbox)
+		return;
+
+	for (i = 0; i < ext_mailbox->count; i++) {
+		id = ext_mailbox->descriptors[i].mailbox->mailbox_id;
+		etdev_dbg(group->etdev, "Disabling mailbox: %d\n", id);
+		if (edgetpu_mailbox_deactivate(group->etdev, id))
+			etdev_err(group->etdev, "Deactivate mailbox %d failed", id);
+	}
+	/*
+	 * 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);
+}
+
+int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, int mailbox_id,
+			       struct edgetpu_external_mailbox_req *ext_mailbox_req)
 {
 	int ret;
 
 	if (!edgetpu_wakelock_lock(client->wakelock)) {
-		etdev_err(client->etdev,
-			  "Enabling mailbox %d needs wakelock acquired\n",
+		etdev_err(client->etdev, "Enabling mailbox %d needs wakelock acquired\n",
 			  mailbox_id);
 		edgetpu_wakelock_unlock(client->wakelock);
 		return -EAGAIN;
 	}
 
-	edgetpu_wakelock_inc_event_locked(client->wakelock,
-					  EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
-
+	if (mailbox_id == EDGETPU_MAILBOX_ID_USE_ASSOC) {
+		ret = edgetpu_mailbox_external_alloc_enable(client, ext_mailbox_req);
+		goto out;
+	}
 	etdev_dbg(client->etdev, "Enabling mailbox: %d\n", mailbox_id);
 
 	ret = edgetpu_mailbox_activate(client->etdev, mailbox_id, -1, false);
 	if (ret)
-		etdev_err(client->etdev, "Activate mailbox %d failed: %d",
-			  mailbox_id, ret);
+		etdev_err(client->etdev, "Activate mailbox %d failed: %d", mailbox_id, ret);
+out:
+	if (!ret)
+		edgetpu_wakelock_inc_event_locked(client->wakelock,
+						  EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
 	edgetpu_wakelock_unlock(client->wakelock);
 	return ret;
 }
 
-int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, u32 mailbox_id)
+int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id)
 {
 	int ret;
 
 	if (!edgetpu_wakelock_lock(client->wakelock)) {
-		etdev_err(client->etdev,
-			  "Disabling mailbox %d needs wakelock acquired\n",
+		etdev_err(client->etdev, "Disabling mailbox %d needs wakelock acquired\n",
 			  mailbox_id);
 		edgetpu_wakelock_unlock(client->wakelock);
 		return -EAGAIN;
 	}
 
-	edgetpu_wakelock_dec_event_locked(client->wakelock,
-					  EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
-
+	if (mailbox_id == EDGETPU_MAILBOX_ID_USE_ASSOC) {
+		ret = edgetpu_mailbox_external_disable_free(client);
+		goto out;
+	}
 	etdev_dbg(client->etdev, "Disabling mailbox: %d\n", mailbox_id);
-	ret = edgetpu_mailbox_deactivate(client->etdev, mailbox_id);
 
+	ret = edgetpu_mailbox_deactivate(client->etdev, mailbox_id);
 	if (ret)
-		etdev_err(client->etdev, "Deactivate mailbox %d failed: %d",
-			  mailbox_id, ret);
+		etdev_err(client->etdev, "Deactivate mailbox %d failed: %d", mailbox_id, ret);
+
+out:
+	if (!ret)
+		edgetpu_wakelock_dec_event_locked(client->wakelock,
+						  EDGETPU_WAKELOCK_EVENT_EXT_MAILBOX);
 	edgetpu_wakelock_unlock(client->wakelock);
 	return ret;
 }
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index d0fb9ae..1ae4889 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -31,6 +31,9 @@
 /* Size of CSRs start from cmd_queue_csr_base can be mmap-ed to userspace. */
 #define USERSPACE_CSR_SIZE 0x1000ul
 
+/* Mailbox ID to indicate external mailboxes */
+#define EDGETPU_MAILBOX_ID_USE_ASSOC -1
+
 struct edgetpu_device_group;
 
 struct edgetpu_mailbox {
@@ -78,6 +81,34 @@ struct edgetpu_vii {
 	edgetpu_queue_mem resp_queue_mem;
 };
 
+/* Structure to hold info about mailbox and its queues. */
+struct edgetpu_mailbox_descriptor {
+	struct edgetpu_mailbox *mailbox;
+	edgetpu_queue_mem cmd_queue_mem;
+	edgetpu_queue_mem resp_queue_mem;
+};
+
+/* 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;
+	/* Leader of device group. */
+	struct edgetpu_dev *etdev;
+	/* Array of external mailboxes info with length @count. */
+	struct edgetpu_mailbox_descriptor *descriptors;
+	/* Mailbox attribute for allocated external mailboxes. */
+	struct edgetpu_mailbox_attr attr;
+};
+
+/* Structure used for requesting to allocate external mailboxes. */
+struct edgetpu_external_mailbox_req {
+	uint start; /* starting index of external mailbox in mailbox_manager */
+	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;
+	struct edgetpu_mailbox_attr attr; /* mailbox attribute for allocation */
+};
+
 /*
  * Structure for recording the driver state vs FW state.
  *
@@ -107,6 +138,8 @@ struct edgetpu_mailbox_manager {
 	u8 vii_index_from, vii_index_to;
 	/* indices reserved for P2P, the range is [from, to) */
 	u8 p2p_index_from, p2p_index_to;
+	/* indices reserved for external mailboxes */
+	u8 ext_index_from, ext_index_to;
 	rwlock_t mailboxes_lock;	/* protects mailboxes */
 	struct edgetpu_mailbox **mailboxes;
 	/* converts index (0 ~ num_mailbox - 1) of mailbox to CSR offset */
@@ -121,6 +154,7 @@ struct edgetpu_mailbox_manager_desc {
 	u8 num_mailbox;
 	u8 num_vii_mailbox;
 	u8 num_p2p_mailbox;
+	u8 num_ext_mailbox;
 	get_csr_base_t get_context_csr_base;
 	get_csr_base_t get_cmd_queue_csr_base;
 	get_csr_base_t get_resp_queue_csr_base;
@@ -302,11 +336,18 @@ void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev);
 int edgetpu_mailbox_p2p_batch(struct edgetpu_mailbox_manager *mgr, uint n,
 			      uint skip_i, struct edgetpu_mailbox **mailboxes);
 
-/* Notify firmware of an external mailbox becoming active */
-int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, u32 mailbox_id);
+/*
+ * If @mailbox_id is EDGETPU_MAILBOX_ID_USE_ASSOC, use @ext_mailbox_req to
+ * allocate external mailboxes and activate the allocated mailboxes.
+ * Otherwise, activate the external mailbox with id @mailbox_id.
+ */
+int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, int mailbox_id,
+			       struct edgetpu_external_mailbox_req *ext_mailbox_req);
 
-/* Notify firmware of an external mailbox becoming inactive */
-int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, u32 mailbox_id);
+/*
+ * Notify firmware of an external mailboxes becoming inactive.
+ */
+int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id);
 
 /*
  * Activates @mailbox_id, OPEN_DEVICE KCI will be sent.
@@ -324,6 +365,12 @@ int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, s16 vcid
 int edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id);
 /* Sets @eh->fw_state to 0. */
 void edgetpu_handshake_clear_fw_state(struct edgetpu_handshake *eh);
+/*
+ * Disables and frees any external mailboxes allocated for @group.
+ *
+ * Caller must hold @group->lock.
+ */
+void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *group);
 
 /* Utilities of circular queue operations */
 
diff --git a/drivers/edgetpu/edgetpu-mapping.c b/drivers/edgetpu/edgetpu-mapping.c
index 6b21db3..136cf29 100644
--- a/drivers/edgetpu/edgetpu-mapping.c
+++ b/drivers/edgetpu/edgetpu-mapping.c
@@ -51,7 +51,7 @@ int edgetpu_mapping_add(struct edgetpu_mapping_root *mappings,
 	struct rb_node *parent = NULL;
 	int ret = -EBUSY;
 
-	if (WARN_ON(!map->release))
+	if (!map->release)
 		return -EINVAL;
 	edgetpu_mapping_lock(mappings);
 	new = &mappings->rb.rb_node;
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index 63fc91c..2b4dd09 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -25,6 +25,7 @@ struct edgetpu_thermal {
 	unsigned long cooling_state;
 	unsigned int tpu_num_states;
 	struct edgetpu_dev *etdev;
+	bool thermal_suspended; /* TPU thermal suspended state */
 };
 
 struct edgetpu_state_pwr {
@@ -40,4 +41,37 @@ struct edgetpu_state_pwr {
 struct edgetpu_thermal *devm_tpu_thermal_create(struct device *dev,
 						struct edgetpu_dev *etdev);
 
+/*
+ * Holds thermal->lock.
+ *
+ * Does nothing if the thermal management is not supported.
+ */
+static inline void edgetpu_thermal_lock(struct edgetpu_thermal *thermal)
+{
+	if (!IS_ERR_OR_NULL(thermal))
+		mutex_lock(&thermal->lock);
+}
+
+/*
+ * Checks whether device is thermal suspended.
+ * Returns false if the thermal management is not supported.
+ */
+static inline bool edgetpu_thermal_is_suspended(struct edgetpu_thermal *thermal)
+{
+	if (!IS_ERR_OR_NULL(thermal))
+		return thermal->thermal_suspended;
+	return false;
+}
+
+/*
+ * Releases thermal->lock.
+ *
+ * Does nothing if the thermal management is not supported.
+ */
+static inline void edgetpu_thermal_unlock(struct edgetpu_thermal *thermal)
+{
+	if (!IS_ERR_OR_NULL(thermal))
+		mutex_unlock(&thermal->lock);
+}
+
 #endif /* __EDGETPU_THERMAL_H__ */
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index c13b42e..0001210 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -13,7 +13,6 @@
 #include "edgetpu-usage-stats.h"
 
 #if IS_ENABLED(CONFIG_ABROLHOS)
-//TODO(b/179343138): Implement for Janeiro
 #include "abrolhos-pm.h"
 
 static enum tpu_pwr_state tpu_states_arr[] = {
diff --git a/drivers/edgetpu/edgetpu.h b/drivers/edgetpu/edgetpu.h
index afa4fed..db6b6b8 100644
--- a/drivers/edgetpu/edgetpu.h
+++ b/drivers/edgetpu/edgetpu.h
@@ -278,7 +278,7 @@ struct edgetpu_sync_ioctl {
 	_IOW(EDGETPU_IOCTL_BASE, 16, struct edgetpu_sync_ioctl)
 
 struct edgetpu_map_dmabuf_ioctl {
-	/* Deprecated; pass 0 to keep compatibility. */
+	/* Ignored. */
 	__u64 offset;
 	/* Ignored; the entire dma-buf is mapped. */
 	__u64 size;
@@ -561,6 +561,8 @@ struct edgetpu_ext_mailbox_ioctl {
 #define EDGETPU_ERROR_HW_NO_ACCESS	0x8
 /* Various hardware failures */
 #define EDGETPU_ERROR_HW_FAIL		0x10
+/* Firmware-reported timeout on runtime processing of workload */
+#define EDGETPU_ERROR_RUNTIME_TIMEOUT	0x20
 
 /*
  * Return fatal errors raised for the client's device group, as a bitmask of