Merge aosp/android-gs-raviole-5.10-android14-qpr2 into aosp/android14-gs-pixel-6.1android14-gs-pixel-6.1

* aosp/android-gs-raviole-5.10-android14-qpr2:
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  [Copybara Auto Merge] Merge branch whitechapel into partner-android
  edgetpu: abrolhos: add pre-allocated iommu domains
  Fix include file name to gs_tmu.h

Change-Id: I9bcd9af94935c0140d6a5f8389b584b35c813bf2
Signed-off-by: Will McVicker <willmcvicker@google.com>

32 files changed, 978 insertions, 473 deletions

diff --git a/drivers/edgetpu/.gitignore b/drivers/edgetpu/.gitignore
new file mode 100644
index 0000000..c04c6ad
--- /dev/null
+++ b/drivers/edgetpu/.gitignore
@@ -0,0 +1 @@
+/gcip-kernel-driver
diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index c7c9cfe..a16c06b 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -11,14 +11,11 @@ else
 	ccflags-y	+= -DGIT_REPO_TAG=\"Not\ a\ git\ repository\"
 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-domain-pool.o
+edgetpu-objs	:= edgetpu-mailbox.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-domain-pool.o
 
-abrolhos-objs	:= abrolhos-core.o abrolhos-debug-dump.o \
-		   abrolhos-device-group.o abrolhos-fs.o abrolhos-device.o \
-		   abrolhos-firmware.o abrolhos-iommu.o \
-		   abrolhos-platform.o abrolhos-pm.o abrolhos-thermal.o \
-		   abrolhos-usage-stats.o abrolhos-wakelock.o \
-		   $(edgetpu-objs)
+mobile-objs	:= edgetpu-kci.o
+
+abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o abrolhos-iommu.o abrolhos-debug-dump.o abrolhos-usage-stats.o abrolhos-wakelock.o $(mobile-objs) $(edgetpu-objs)
 
 # This -I is for the trace file include. It was removed from ccflags-y to avoid
 # including the header stubs.
diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index 0aecdd1..bc1108f 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -93,18 +93,12 @@ void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
 				     struct edgetpu_kci_response_element *resp)
 {
 	switch (resp->code) {
-	case RKCI_CODE_PM_QOS:
-		mobile_pm_set_pm_qos(etdev, resp->retval);
-		break;
-	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);
+			edgetpu_mobile_pm_set_pm_qos(etdev, resp->retval);
 		if (resp->status != (typeof(resp->status))~0ull)
-			mobile_pm_set_bts(etdev, resp->status);
+			edgetpu_mobile_pm_set_bts(etdev, resp->status);
 		break;
 	default:
 		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n",
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 8bf40e2..015e755 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -85,5 +85,5 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 	platform_pwr->after_create = abrolhos_pm_after_create;
 	platform_pwr->acpm_set_rate = exynos_acpm_set_rate;
 
-	return mobile_pm_create(etdev);
+	return edgetpu_mobile_pm_create(etdev);
 }
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index 710adbf..883bdc5 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -1,4 +1,4 @@
 // SPDX-License-Identifier: GPL-2.0
 #include <linux/thermal.h>
-#include <soc/google/gs101_tmu.h>
+#include <soc/google/gs_tmu.h>
 #include "mobile-thermal.c"
diff --git a/drivers/edgetpu/abrolhos/config.h b/drivers/edgetpu/abrolhos/config.h
index 57d6170..b1db928 100644
--- a/drivers/edgetpu/abrolhos/config.h
+++ b/drivers/edgetpu/abrolhos/config.h
@@ -20,6 +20,9 @@
 /* Reserved VCID that uses the extra partition. */
 #define EDGETPU_VCID_EXTRA_PARTITION 0
 
+/* Pre-allocate 1 IOMMU domain per VCID */
+#define EDGETPU_NUM_PREALLOCATED_DOMAINS EDGETPU_NUM_VCIDS
+
 /* Is a "mobile" style device. */
 #define EDGETPU_FEATURE_MOBILE
 #define EDGETPU_HAS_WAKELOCK
@@ -31,6 +34,9 @@
  */
 #define EDGETPU_HAS_REMAPPED_DATA
 
+/* Metrics are reported for a single default "cluster" component. */
+#define EDGETPU_TPU_CLUSTER_COUNT	1
+
 /*
  * The TPU VA where the firmware is located.
  *
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index ff1ad77..8d66b81 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -24,4 +24,10 @@
 #define EDGETPU_NUM_CORES 1
 #endif
 
+/* Uses a smaller size for unittests to avoid DMA warnings. */
+#if IS_ENABLED(CONFIG_EDGETPU_TEST)
+#undef EDGETPU_DEBUG_DUMP_MEM_SIZE
+#define EDGETPU_DEBUG_DUMP_MEM_SIZE 0x20000
+#endif
+
 #endif /* __EDGETPU_CONFIG_H__ */
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 49f19a1..3cce9c1 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -490,6 +490,8 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 		etdev_warn(etdev, "debug dump init fail: %d", ret);
 
 	edgetpu_chip_init(etdev);
+	/* No limit on DMA segment size */
+	dma_set_max_seg_size(etdev->dev, UINT_MAX);
 	return 0;
 
 remove_kci:
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index c5b6647..c410ebc 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -27,6 +27,7 @@
 #include "edgetpu-async.h"
 #include "edgetpu-config.h"
 #include "edgetpu-device-group.h"
+#include "edgetpu-dmabuf.h"
 #include "edgetpu-dram.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -140,7 +141,8 @@ static int edgetpu_group_activate(struct edgetpu_device_group *group)
 		return 0;
 
 	mailbox_id = edgetpu_group_context_id_locked(group);
-	ret = edgetpu_mailbox_activate(group->etdev, mailbox_id, group->vcid, !group->activated);
+	ret = edgetpu_mailbox_activate(group->etdev, mailbox_id, group->mbox_attr.client_priv,
+				       group->vcid, !group->activated);
 	if (ret) {
 		etdev_err(group->etdev, "activate mailbox for VCID %d failed with %d", group->vcid,
 			  ret);
@@ -505,6 +507,8 @@ static void edgetpu_device_group_release(struct edgetpu_device_group *group)
 		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
 		edgetpu_mmu_free_domain(group->etdev, group->etdomain);
 	}
+	/* Signal any unsignaled dma fences owned by the group with an error. */
+	edgetpu_sync_fence_group_shutdown(group);
 	group->status = EDGETPU_DEVICE_GROUP_DISBANDED;
 }
 
@@ -721,6 +725,7 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 	group->vii.etdev = client->etdev;
 	mutex_init(&group->lock);
 	rwlock_init(&group->events.lock);
+	INIT_LIST_HEAD(&group->dma_fence_list);
 	edgetpu_mapping_init(&group->host_mappings);
 	edgetpu_mapping_init(&group->dmabuf_mappings);
 	group->mbox_attr = *attr;
@@ -1174,12 +1179,9 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 		return ERR_PTR(-EFAULT);
 	}
 	offset = host_addr & (PAGE_SIZE - 1);
-	/* overflow check (should also be caught by access_ok) */
-	if (unlikely((size + offset) / PAGE_SIZE >= UINT_MAX - 1 || size + offset < size)) {
-		etdev_err(etdev, "address overflow in buffer map request");
-		return ERR_PTR(-EFAULT);
-	}
 	num_pages = DIV_ROUND_UP((size + offset), PAGE_SIZE);
+	if (num_pages * PAGE_SIZE < size + offset)
+		return ERR_PTR(-EINVAL);
 	etdev_dbg(etdev, "%s: hostaddr=%#llx pages=%u", __func__, host_addr, num_pages);
 	/*
 	 * "num_pages" is decided from user-space arguments, don't show warnings
@@ -1196,6 +1198,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	 * it with FOLL_WRITE.
 	 * default to read/write if find_extend_vma returns NULL
 	 */
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	down_read(&current->mm->mmap_sem);
+#else
+	mmap_read_lock(current->mm);
+#endif
 	vma = find_extend_vma(current->mm, host_addr & PAGE_MASK);
 	if (vma && !(vma->vm_flags & VM_WRITE)) {
 		foll_flags &= ~FOLL_WRITE;
@@ -1203,6 +1210,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	} else {
 		*preadonly = false;
 	}
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	up_read(&current->mm->mmap_sem);
+#else
+	mmap_read_unlock(current->mm);
+#endif
 
 	/* Try fast call first, in case it's actually faster. */
 	ret = pin_user_pages_fast(host_addr & PAGE_MASK, num_pages, foll_flags,
@@ -1246,8 +1258,18 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 		kvfree(pages);
 		return ERR_PTR(-ENOMEM);
 	}
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	down_read(&current->mm->mmap_sem);
+#else
+	mmap_read_lock(current->mm);
+#endif
 	ret = pin_user_pages(host_addr & PAGE_MASK, num_pages, foll_flags,
 			     pages, vmas);
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	up_read(&current->mm->mmap_sem);
+#else
+	mmap_read_unlock(current->mm);
+#endif
 	kvfree(vmas);
 	if (ret < 0) {
 		etdev_dbg(etdev, "pin_user_pages failed %u:%pK-%u: %d",
@@ -1265,6 +1287,8 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 			  "pin_user_pages partial %u:%pK npages=%u pinned=%d",
 			  group->workload_id, (void *)host_addr, num_pages,
 			  ret);
+		etdev_err(etdev, "can only lock %u of %u pages requested",
+			  (unsigned int)ret, num_pages);
 		num_pages = ret;
 		ret = -EFAULT;
 		goto error;
@@ -1866,7 +1890,7 @@ uint edgetpu_group_get_fatal_errors(struct edgetpu_device_group *group)
 	uint fatal_errors;
 
 	mutex_lock(&group->lock);
-	fatal_errors = group->fatal_errors;
+	fatal_errors = edgetpu_group_get_fatal_errors_locked(group);
 	mutex_unlock(&group->lock);
 	return fatal_errors;
 }
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index bf2776b..a8a3bac 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -129,6 +129,9 @@ struct edgetpu_device_group {
 	/* Mask of errors set for this group. */
 	uint fatal_errors;
 
+	/* List of DMA fences owned by this group */
+	struct list_head dma_fence_list;
+
 	/* end of fields protected by @lock */
 
 	/* TPU IOVA mapped to host DRAM space */
@@ -207,15 +210,27 @@ edgetpu_device_group_is_disbanded(const struct edgetpu_device_group *group)
 }
 
 /*
- * Returns -ECANCELED if the status of group is ERRORED, otherwise returns
- * -EINVAL.
+ * Return fatal error status for the group.
+ *
+ * Caller holds @group->lock.
+ */
+static inline uint edgetpu_group_get_fatal_errors_locked(struct edgetpu_device_group *group)
+{
+	return group->fatal_errors;
+}
+
+/*
+ * Returns -ECANCELED if the status of group is ERRORED, otherwise returns -EINVAL.
  *
  * Caller holds @group->lock.
  */
 static inline int edgetpu_group_errno(struct edgetpu_device_group *group)
 {
-	if (edgetpu_device_group_is_errored(group))
+	if (edgetpu_device_group_is_errored(group)) {
+		etdev_err(group->etdev, "group %u error status 0x%x\n", group->workload_id,
+			  edgetpu_group_get_fatal_errors_locked(group));
 		return -ECANCELED;
+	}
 	return -EINVAL;
 }
 
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 657ae75..fbc9f48 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -72,6 +72,9 @@ struct edgetpu_dmabuf_map {
  * @fence:		the base DMA fence
  * @lock:		spinlock protecting updates to @fence
  * @timeline_name:	name of the timeline associated with the fence
+ * @group:		owning device group
+ * @etfence_list:	global list of all edgetpu DMA fences
+ * @group_list:		list of DMA fences owned by the same group
  *
  * It is likely timelines will become a separate object in the future,
  * but for now there's a unique named timeline associated with each fence.
@@ -80,7 +83,9 @@ struct edgetpu_dma_fence {
 	struct dma_fence fence;
 	spinlock_t lock;
 	char timeline_name[EDGETPU_SYNC_TIMELINE_NAME_LEN];
+	struct edgetpu_device_group *group;
 	struct list_head etfence_list;
+	struct list_head group_list;
 };
 
 /* List of all edgetpu fence objects for debugging. */
@@ -873,6 +878,7 @@ static const char *edgetpu_dma_fence_get_timeline_name(struct dma_fence *fence)
 static void edgetpu_dma_fence_release(struct dma_fence *fence)
 {
 	struct edgetpu_dma_fence *etfence = to_etfence(fence);
+	struct edgetpu_device_group *group;
 	unsigned long flags;
 
 	if (!etfence)
@@ -881,6 +887,17 @@ static void edgetpu_dma_fence_release(struct dma_fence *fence)
 	spin_lock_irqsave(&etfence_list_lock, flags);
 	list_del(&etfence->etfence_list);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
+
+	/* group might not yet be set if error at init time. */
+	group = etfence->group;
+	if (group) {
+		mutex_lock(&group->lock);
+		list_del(&etfence->group_list);
+		mutex_unlock(&group->lock);
+		/* Release this fence's reference on the owning group. */
+		edgetpu_device_group_put(group);
+	}
+
 	kfree(etfence);
 }
 
@@ -904,7 +921,8 @@ static const struct dma_fence_ops edgetpu_dma_fence_ops = {
 #define SEQ_FMT "%llu"
 #endif
 
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap)
 {
 	int fd = get_unused_fd_flags(O_CLOEXEC);
 	int ret;
@@ -926,6 +944,7 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	 * list_head is needed for list_del().
 	 */
 	INIT_LIST_HEAD(&etfence->etfence_list);
+	INIT_LIST_HEAD(&etfence->group_list);
 	memcpy(&etfence->timeline_name, &datap->timeline_name,
 	       EDGETPU_SYNC_TIMELINE_NAME_LEN - 1);
 
@@ -944,7 +963,10 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	spin_lock_irqsave(&etfence_list_lock, flags);
 	list_add_tail(&etfence->etfence_list, &etfence_list_head);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
-
+	etfence->group = edgetpu_device_group_get(group);
+	mutex_lock(&group->lock);
+	list_add_tail(&etfence->group_list, &group->dma_fence_list);
+	mutex_unlock(&group->lock);
 	fd_install(fd, sync_file->file);
 	datap->fence = fd;
 	return 0;
@@ -954,26 +976,14 @@ err_put_fd:
 	return ret;
 }
 
-int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+static int _edgetpu_sync_fence_signal(struct dma_fence *fence, int errno, bool ignore_signaled)
 {
-	struct dma_fence *fence;
-	int errno;
 	int ret;
 
-	errno = datap->error;
-	if (errno > 0)
-		errno = -errno;
-	if (errno < -MAX_ERRNO)
-		return -EINVAL;
-
-	fence = sync_file_get_fence(datap->fence);
-	if (!fence)
-		return -EINVAL;
-
 	spin_lock_irq(fence->lock);
 	/* don't signal fence twice */
 	if (unlikely(test_bit(DMA_FENCE_FLAG_SIGNALED_BIT, &fence->flags))) {
-		ret = -EINVAL;
+		ret = ignore_signaled ? 0 : -EINVAL;
 		goto out_unlock;
 	}
 	pr_debug("%s: %s-%s%llu-" SEQ_FMT " errno=%d\n", __func__,
@@ -986,10 +996,51 @@ int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
 
 out_unlock:
 	spin_unlock_irq(fence->lock);
+	return ret;
+}
+
+int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+{
+	struct dma_fence *fence;
+	int errno;
+	int ret;
+
+	errno = datap->error;
+	if (errno > 0)
+		errno = -errno;
+	if (errno < -MAX_ERRNO)
+		return -EINVAL;
+
+	fence = sync_file_get_fence(datap->fence);
+	if (!fence)
+		return -EINVAL;
+
+	ret = _edgetpu_sync_fence_signal(fence, errno, false);
 	dma_fence_put(fence);
 	return ret;
 }
 
+/* Caller holds group lock. */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group)
+{
+	struct list_head *pos;
+	int ret;
+
+	lockdep_assert_held(&group->lock);
+	list_for_each(pos, &group->dma_fence_list) {
+		struct edgetpu_dma_fence *etfence =
+			container_of(pos, struct edgetpu_dma_fence, group_list);
+		struct dma_fence *fence = &etfence->fence;
+
+		ret = _edgetpu_sync_fence_signal(fence, -EPIPE, true);
+		if (ret)
+			etdev_warn(group->etdev, "error %d signaling fence %s-%s %llu-" SEQ_FMT,
+				   ret, fence->ops->get_driver_name(fence),
+				   fence->ops->get_timeline_name(fence),
+				   fence->context, fence->seqno);
+	}
+}
+
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap)
 {
 	struct dma_fence *fence;
@@ -1027,8 +1078,9 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 
 		spin_lock_irq(&etfence->lock);
 		seq_printf(s, "%s-%s %llu-" SEQ_FMT " %s",
-			   edgetpu_dma_fence_get_driver_name(fence),
-			   etfence->timeline_name, fence->context, fence->seqno,
+			   fence->ops->get_driver_name(fence),
+			   fence->ops->get_timeline_name(fence),
+			   fence->context, fence->seqno,
 			   sync_status_str(dma_fence_get_status_locked(fence)));
 
 		if (test_bit(DMA_FENCE_FLAG_TIMESTAMP_BIT, &fence->flags)) {
@@ -1041,8 +1093,7 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 
 		if (fence->error)
 			seq_printf(s, " err=%d", fence->error);
-
-		seq_putc(s, '\n');
+		seq_printf(s, " group=%u\n", etfence->group->workload_id);
 		spin_unlock_irq(&etfence->lock);
 	}
 
diff --git a/drivers/edgetpu/edgetpu-dmabuf.h b/drivers/edgetpu/edgetpu-dmabuf.h
index f35b2d3..2b9e281 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.h
+++ b/drivers/edgetpu/edgetpu-dmabuf.h
@@ -38,11 +38,17 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 int edgetpu_unmap_bulk_dmabuf(struct edgetpu_device_group *group,
 			      tpu_addr_t tpu_addr);
 /* Create a DMA sync fence via ioctl */
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap);
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap);
 /* Signal a DMA sync fence, optionally specifying error status */
 int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap);
 /* Return DMA sync fence status */
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap);
+/*
+ * Send error signal to any remaining unsignalled DMA sync fences in a group being disbanded.
+ * Caller holds group lock.
+ */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group);
 /* Dump sync fence info from debugfs */
 int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused);
 
diff --git a/drivers/edgetpu/edgetpu-external.c b/drivers/edgetpu/edgetpu-external.c
index 4b86e13..c983d3c 100644
--- a/drivers/edgetpu/edgetpu-external.c
+++ b/drivers/edgetpu/edgetpu-external.c
@@ -95,12 +95,13 @@ static int edgetpu_external_mailbox_alloc(struct device *edgetpu_dev,
 
 	if (copy_from_user(&req.attr, (void __user *)client_info->attr, sizeof(req.attr))) {
 		if (!client_info->attr)
-			etdev_warn(client->etdev,
-				   "Illegal mailbox attributes, using VII mailbox attrs\n");
+			etdev_dbg(client->etdev,
+				  "Using VII mailbox attrs for external mailbox\n");
 		req.attr = group->mbox_attr;
 	}
 
-	ret = edgetpu_mailbox_enable_ext(client, EDGETPU_MAILBOX_ID_USE_ASSOC, &req);
+	ret = edgetpu_mailbox_enable_ext(client, EDGETPU_MAILBOX_ID_USE_ASSOC, &req,
+					 group->mbox_attr.client_priv);
 	if (ret)
 		goto error_put_group;
 	mutex_lock(&group->lock);
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 1ef1354..ad27ec9 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -5,6 +5,7 @@
  * Copyright (C) 2019-2020 Google, Inc.
  */
 
+#include <linux/debugfs.h>
 #include <linux/delay.h>
 #include <linux/device.h>
 #include <linux/firmware.h>
@@ -29,6 +30,30 @@
 static char *firmware_name;
 module_param(firmware_name, charp, 0660);
 
+/*
+ * Any tracing level vote with the following bit set will be considered as a default vote.
+ */
+#define EDGETPU_FW_TRACING_DEFAULT_VOTE BIT(8)
+
+struct edgetpu_fw_tracing {
+	struct device *dev;
+	struct dentry *dentry;
+
+	/*
+	 * Lock to protect the struct members listed below.
+	 *
+	 * Note that since the request of tracing level adjusting might happen during power state
+	 * transitions (i.e., another thread calling edgetpu_firmware_tracing_restore_on_powering()
+	 * with pm lock held), one must either use the non-blocking edgetpu_pm_trylock() or make
+	 * sure there won't be any new power transition after holding this lock to prevent deadlock.
+	 */
+	struct mutex lock;
+	/* Actual firmware tracing level. */
+	unsigned long active_level;
+	/* Requested firmware tracing level. */
+	unsigned long request_level;
+};
+
 struct edgetpu_firmware_private {
 	const struct edgetpu_firmware_chip_data *chip_fw;
 	void *data; /* for edgetpu_firmware_(set/get)_data */
@@ -38,6 +63,7 @@ struct edgetpu_firmware_private {
 	struct edgetpu_firmware_desc bl1_fw_desc;
 	enum edgetpu_firmware_status status;
 	struct edgetpu_fw_info fw_info;
+	struct edgetpu_fw_tracing fw_tracing;
 };
 
 void edgetpu_firmware_set_data(struct edgetpu_firmware *et_fw, void *data)
@@ -134,6 +160,124 @@ static char *fw_flavor_str(enum edgetpu_fw_flavor fw_flavor)
 	return "?";
 }
 
+static int edgetpu_firmware_tracing_active_get(void *data, u64 *val)
+{
+	struct edgetpu_firmware *et_fw = data;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+
+	mutex_lock(&fw_tracing->lock);
+	*val = fw_tracing->active_level;
+	mutex_unlock(&fw_tracing->lock);
+
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_edgetpu_firmware_tracing_active, edgetpu_firmware_tracing_active_get,
+			 NULL, "%llu\n");
+
+static int edgetpu_firmware_tracing_request_get(void *data, u64 *val)
+{
+	struct edgetpu_firmware *et_fw = data;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+
+	mutex_lock(&fw_tracing->lock);
+	*val = fw_tracing->request_level;
+	mutex_unlock(&fw_tracing->lock);
+
+	return 0;
+}
+
+/*
+ * fw_tracing->lock may optionally be held if the caller wants the new level to be set as a
+ * critical section.  If not held the caller is syncing current tracing level but not as a critical
+ * section with the calling code.  Firmware tracing levels are not expected to change frequently or
+ * via concurrent requests.  Only the code that restore the tracing level at power up requires
+ * consistency with the state managed by the calling code.  Since this code is called as part of
+ * power up processing, in order to avoid deadlocks, most callers set a requested state and then
+ * sync the current state to firmware (if powered on) without holding the lock across the powered-on
+ * check, with no harm done if the requested state changed again using a concurrent request.
+ */
+static int edgetpu_firmware_tracing_set_level(struct edgetpu_firmware *et_fw)
+{
+	unsigned long active_level;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+	int ret = edgetpu_kci_firmware_tracing_level(etdev, fw_tracing->request_level,
+						     &active_level);
+
+	if (ret)
+		etdev_warn(et_fw->etdev, "Failed to set firmware tracing level to %lu: %d",
+			   fw_tracing->request_level, ret);
+	else
+		fw_tracing->active_level =
+			(fw_tracing->request_level & EDGETPU_FW_TRACING_DEFAULT_VOTE) ?
+				EDGETPU_FW_TRACING_DEFAULT_VOTE : active_level;
+
+	return ret;
+}
+
+static int edgetpu_firmware_tracing_request_set(void *data, u64 val)
+{
+	struct edgetpu_firmware *et_fw = data;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+	int ret = 0;
+
+	mutex_lock(&fw_tracing->lock);
+	fw_tracing->request_level = val;
+	mutex_unlock(&fw_tracing->lock);
+
+	if (edgetpu_pm_get_if_powered(etdev->pm)) {
+		ret = edgetpu_firmware_tracing_set_level(et_fw);
+		edgetpu_pm_put(etdev->pm);
+	}
+
+	return ret;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_edgetpu_firmware_tracing_request,
+			 edgetpu_firmware_tracing_request_get, edgetpu_firmware_tracing_request_set,
+			 "%llu\n");
+
+static void edgetpu_firmware_tracing_init(struct edgetpu_firmware *et_fw)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+
+	fw_tracing->active_level = EDGETPU_FW_TRACING_DEFAULT_VOTE;
+	fw_tracing->request_level = EDGETPU_FW_TRACING_DEFAULT_VOTE;
+	mutex_init(&fw_tracing->lock);
+
+	fw_tracing->dentry = debugfs_create_dir("fw_tracing", etdev->d_entry);
+	if (IS_ERR(fw_tracing->dentry)) {
+		etdev_warn(etdev, "Failed to create fw tracing debugfs interface");
+		return;
+	}
+
+	debugfs_create_file("active", 0440, fw_tracing->dentry, et_fw,
+			    &fops_edgetpu_firmware_tracing_active);
+	debugfs_create_file("request", 0660, fw_tracing->dentry, et_fw,
+			    &fops_edgetpu_firmware_tracing_request);
+}
+
+static void edgetpu_firmware_tracing_destroy(struct edgetpu_firmware *et_fw)
+{
+	debugfs_remove_recursive(et_fw->p->fw_tracing.dentry);
+}
+
+static int edgetpu_firmware_tracing_restore_on_powering(struct edgetpu_firmware *et_fw)
+{
+	int ret = 0;
+	struct edgetpu_fw_tracing *fw_tracing = &et_fw->p->fw_tracing;
+
+	mutex_lock(&fw_tracing->lock);
+	fw_tracing->active_level = EDGETPU_FW_TRACING_DEFAULT_VOTE;
+	if (!(fw_tracing->request_level & EDGETPU_FW_TRACING_DEFAULT_VOTE))
+		ret = edgetpu_firmware_tracing_set_level(et_fw);
+	mutex_unlock(&fw_tracing->lock);
+	return ret;
+}
+
 static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;
@@ -172,6 +316,9 @@ static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 
 		if (ret)
 			etdev_warn(etdev, "telemetry KCI error: %d", ret);
+		ret = edgetpu_firmware_tracing_restore_on_powering(et_fw);
+		if (ret)
+			etdev_warn_ratelimited(etdev, "firmware tracing restore error: %d", ret);
 		/* Set debug dump buffer in FW */
 		edgetpu_get_debug_dump(etdev, 0);
 	}
@@ -331,20 +478,21 @@ int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 				enum edgetpu_firmware_flags flags)
 {
 	const struct edgetpu_firmware_chip_data *chip_fw = et_fw->p->chip_fw;
+	struct edgetpu_dev *etdev = et_fw->etdev;
 	struct edgetpu_firmware_desc new_fw_desc;
 	int ret;
 	bool is_bl1_run = (flags & FW_BL1);
 
 	edgetpu_firmware_set_loading(et_fw);
 	if (!is_bl1_run)
-		edgetpu_sw_wdt_stop(et_fw->etdev);
+		edgetpu_sw_wdt_stop(etdev);
 
 	memset(&new_fw_desc, 0, sizeof(new_fw_desc));
 	ret = edgetpu_firmware_load_locked(et_fw, &new_fw_desc, name, flags);
 	if (ret)
 		goto out_failed;
 
-	etdev_dbg(et_fw->etdev, "run fw %s flags=%#x", name, flags);
+	etdev_dbg(etdev, "run fw %s flags=%#x", name, flags);
 	if (chip_fw->prepare_run) {
 		/* Note this may recursively call us to run BL1 */
 		ret = chip_fw->prepare_run(et_fw, &new_fw_desc.buf);
@@ -369,13 +517,16 @@ int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 
 	/* Don't start wdt if loaded firmware is second stage bootloader. */
 	if (!ret && !is_bl1_run && et_fw->p->fw_info.fw_flavor != FW_FLAVOR_BL1)
-		edgetpu_sw_wdt_start(et_fw->etdev);
+		edgetpu_sw_wdt_start(etdev);
 
 	if (!ret && !is_bl1_run && chip_fw->launch_complete)
 		chip_fw->launch_complete(et_fw);
 	else if (ret && chip_fw->launch_failed)
 		chip_fw->launch_failed(et_fw, ret);
 	edgetpu_firmware_set_state(et_fw, ret);
+	/* If previous firmware was metrics v1-only reset that flag and probe this again. */
+	if (etdev->usage_stats)
+		etdev->usage_stats->use_metrics_v1 = false;
 	return ret;
 
 out_unload_new_fw:
@@ -687,6 +838,7 @@ int edgetpu_firmware_create(struct edgetpu_dev *etdev,
 	else
 		edgetpu_sw_wdt_set_handler(
 			etdev, edgetpu_firmware_wdt_timeout_action, etdev);
+	edgetpu_firmware_tracing_init(et_fw);
 	return 0;
 
 out_device_remove_group:
@@ -724,6 +876,7 @@ void edgetpu_firmware_destroy(struct edgetpu_dev *etdev)
 		edgetpu_firmware_unload_locked(et_fw, &et_fw->p->fw_desc);
 		edgetpu_firmware_unload_locked(et_fw, &et_fw->p->bl1_fw_desc);
 		mutex_unlock(&et_fw->p->fw_desc_lock);
+		edgetpu_firmware_tracing_destroy(et_fw);
 	}
 
 	etdev->firmware = NULL;
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index d0efb67..db55484 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -60,18 +60,19 @@ static struct dentry *edgetpu_debugfs_dir;
 #define LOCK(client) mutex_lock(&client->group_lock)
 #define UNLOCK(client) mutex_unlock(&client->group_lock)
 /*
- * Locks @client->group_lock and assigns @client->group to @grp.
- * Returns -EINVAL if @client is not the leader of the group.
+ * Locks @client->group_lock and checks whether @client is the leader of a group.
+ * If @client is not the leader of a group, unlocks group_lock and returns false.
+ * If @client is the leader of a group, returns true with group_lock held.
  */
-#define LOCK_RETURN_IF_NOT_LEADER(client, grp)                                 \
-	do {                                                                   \
-		LOCK(client);                                                  \
-		grp = client->group;                                           \
-		if (!grp || !edgetpu_device_group_is_leader(grp, client)) {    \
-			UNLOCK(client);                                        \
-			return -EINVAL;                                        \
-		}                                                              \
-	} while (0)
+static inline bool lock_check_group_leader(struct edgetpu_client *client)
+{
+	LOCK(client);
+	if (!client->group || !edgetpu_device_group_is_leader(client->group, client)) {
+		UNLOCK(client);
+		return false;
+	}
+	return true;
+}
 
 bool is_edgetpu_file(struct file *file)
 {
@@ -139,16 +140,15 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 static int edgetpu_ioctl_set_eventfd(struct edgetpu_client *client,
 				     struct edgetpu_event_register __user *argp)
 {
-	struct edgetpu_device_group *group;
 	int ret;
 	struct edgetpu_event_register eventreg;
 
 	if (copy_from_user(&eventreg, argp, sizeof(eventreg)))
 		return -EFAULT;
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	ret = edgetpu_group_set_eventfd(group, eventreg.event_id,
-					eventreg.eventfd);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	ret = edgetpu_group_set_eventfd(client->group, eventreg.event_id, eventreg.eventfd);
 	UNLOCK(client);
 	return ret;
 }
@@ -156,10 +156,9 @@ static int edgetpu_ioctl_set_eventfd(struct edgetpu_client *client,
 static int edgetpu_ioctl_unset_eventfd(struct edgetpu_client *client,
 				       uint event_id)
 {
-	struct edgetpu_device_group *group;
-
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	edgetpu_group_unset_eventfd(group, event_id);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	edgetpu_group_unset_eventfd(client->group, event_id);
 	UNLOCK(client);
 	return 0;
 }
@@ -311,9 +310,10 @@ static int edgetpu_ioctl_map_buffer(struct edgetpu_client *client,
 
 	trace_edgetpu_map_buffer_start(&ibuf);
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
 	/* to prevent group being released when we perform map/unmap later */
-	group = edgetpu_device_group_get(group);
+	group = edgetpu_device_group_get(client->group);
 	/*
 	 * Don't hold @client->group_lock on purpose since
 	 * 1. We don't care whether @client still belongs to @group.
@@ -343,15 +343,15 @@ static int edgetpu_ioctl_unmap_buffer(struct edgetpu_client *client,
 				      struct edgetpu_map_ioctl __user *argp)
 {
 	struct edgetpu_map_ioctl ibuf;
-	struct edgetpu_device_group *group;
 	int ret;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	ret = edgetpu_device_group_unmap(group, ibuf.die_index,
-					 ibuf.device_address, ibuf.flags);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	ret = edgetpu_device_group_unmap(client->group, ibuf.die_index, ibuf.device_address,
+					 ibuf.flags);
 	UNLOCK(client);
 	return ret;
 }
@@ -369,14 +369,14 @@ edgetpu_ioctl_allocate_device_buffer(struct edgetpu_client *client, u64 size)
 static int edgetpu_ioctl_sync_buffer(struct edgetpu_client *client,
 				     struct edgetpu_sync_ioctl __user *argp)
 {
-	struct edgetpu_device_group *group;
 	int ret;
 	struct edgetpu_sync_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	ret = edgetpu_device_group_sync_buffer(group, &ibuf);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	ret = edgetpu_device_group_sync_buffer(client->group, &ibuf);
 	UNLOCK(client);
 	return ret;
 }
@@ -394,9 +394,10 @@ edgetpu_ioctl_map_dmabuf(struct edgetpu_client *client,
 
 	trace_edgetpu_map_dmabuf_start(&ibuf);
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
 	/* to prevent group being released when we perform unmap on fault */
-	group = edgetpu_device_group_get(group);
+	group = edgetpu_device_group_get(client->group);
 	ret = edgetpu_map_dmabuf(group, &ibuf);
 	UNLOCK(client);
 	if (ret)
@@ -419,19 +420,20 @@ static int
 edgetpu_ioctl_unmap_dmabuf(struct edgetpu_client *client,
 			   struct edgetpu_map_dmabuf_ioctl __user *argp)
 {
-	struct edgetpu_device_group *group;
 	int ret;
 	struct edgetpu_map_dmabuf_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	ret = edgetpu_unmap_dmabuf(group, ibuf.die_index, ibuf.device_address);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	ret = edgetpu_unmap_dmabuf(client->group, ibuf.die_index, ibuf.device_address);
 	UNLOCK(client);
 	return ret;
 }
 
 static int edgetpu_ioctl_sync_fence_create(
+	struct edgetpu_client *client,
 	struct edgetpu_create_sync_fence_data __user *datap)
 {
 	struct edgetpu_create_sync_fence_data data;
@@ -439,7 +441,14 @@ static int edgetpu_ioctl_sync_fence_create(
 
 	if (copy_from_user(&data, (void __user *)datap, sizeof(data)))
 		return -EFAULT;
-	ret = edgetpu_sync_fence_create(&data);
+	LOCK(client);
+	if (!client->group) {
+		etdev_err(client->etdev, "client creating sync fence not joined to a device group");
+		UNLOCK(client);
+		return -EINVAL;
+	}
+	ret = edgetpu_sync_fence_create(client->group, &data);
+	UNLOCK(client);
 	if (ret)
 		return ret;
 	if (copy_to_user((void __user *)datap, &data, sizeof(data)))
@@ -468,9 +477,10 @@ edgetpu_ioctl_map_bulk_dmabuf(struct edgetpu_client *client,
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
 	/* to prevent group being released when we perform unmap on fault */
-	group = edgetpu_device_group_get(group);
+	group = edgetpu_device_group_get(client->group);
 	ret = edgetpu_map_bulk_dmabuf(group, &ibuf);
 	UNLOCK(client);
 	if (ret)
@@ -489,14 +499,14 @@ static int edgetpu_ioctl_unmap_bulk_dmabuf(
 	struct edgetpu_client *client,
 	struct edgetpu_map_bulk_dmabuf_ioctl __user *argp)
 {
-	struct edgetpu_device_group *group;
 	int ret;
 	struct edgetpu_map_bulk_dmabuf_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
-	ret = edgetpu_unmap_bulk_dmabuf(group, ibuf.device_address);
+	if (!lock_check_group_leader(client))
+		return -EINVAL;
+	ret = edgetpu_unmap_bulk_dmabuf(client->group, ibuf.device_address);
 	UNLOCK(client);
 	return ret;
 }
@@ -590,55 +600,52 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 	 */
 	client->pid = current->pid;
 	client->tgid = current->tgid;
-	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_unlock;
+		       "wakelock acquire rejected due to device thermal limit exceeded");
+		goto error_client_unlock;
 	} 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);
-			goto error_unlock;
+			goto error_client_unlock;
 		}
 	}
 	edgetpu_wakelock_lock(client->wakelock);
 	/* when NO_WAKELOCK: count should be 1 so here is a no-op */
 	count = edgetpu_wakelock_acquire(client->wakelock);
 	if (count < 0) {
-		edgetpu_pm_put(client->etdev->pm);
 		ret = count;
-		goto error_unlock;
+		goto error_wakelock_unlock;
 	}
 	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);
+			etdev_warn(client->etdev, "failed to attach mailbox: %d", ret);
 			edgetpu_wakelock_release(client->wakelock);
-			edgetpu_wakelock_unlock(client->wakelock);
-			goto error_unlock;
+			goto error_wakelock_unlock;
 		}
-	} else {
-		/* Balance the power up count due to pm_get above.*/
-		edgetpu_pm_put(client->etdev->pm);
 	}
+
+error_wakelock_unlock:
 	edgetpu_wakelock_unlock(client->wakelock);
+
+	/* Balance the power up count due to pm_get above.*/
+	if (ret || count)
+		edgetpu_pm_put(client->etdev->pm);
+
+error_client_unlock:
 	UNLOCK(client);
-	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
-		  count + 1);
-	return 0;
-error_unlock:
-	UNLOCK(client);
-	etdev_err(client->etdev, "client pid %d failed to acquire wakelock",
-		  client->pid);
+
+	if (ret)
+		etdev_err(client->etdev, "client pid %d failed to acquire wakelock", client->pid);
+	else
+		etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__, count + 1);
+
 	return ret;
 }
 
@@ -786,7 +793,7 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 		ret = edgetpu_ioctl_allocate_device_buffer(client, (u64)argp);
 		break;
 	case EDGETPU_CREATE_SYNC_FENCE:
-		ret = edgetpu_ioctl_sync_fence_create(argp);
+		ret = edgetpu_ioctl_sync_fence_create(client, argp);
 		break;
 	case EDGETPU_SIGNAL_SYNC_FENCE:
 		ret = edgetpu_ioctl_sync_fence_signal(argp);
@@ -1023,6 +1030,33 @@ static const struct file_operations mappings_ops = {
 	.release = single_release,
 };
 
+static int syncfences_open(struct inode *inode, struct file *file)
+{
+	return single_open(file, edgetpu_sync_fence_debugfs_show, inode->i_private);
+}
+
+static const struct file_operations syncfences_ops = {
+	.open = syncfences_open,
+	.read = seq_read,
+	.llseek = seq_lseek,
+	.owner = THIS_MODULE,
+	.release = single_release,
+};
+
+static int edgetpu_pm_debugfs_set_wakelock(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = data;
+	int ret = 0;
+
+	if (val)
+		ret = edgetpu_pm_get(etdev->pm);
+	else
+		edgetpu_pm_put(etdev->pm);
+	return ret;
+}
+DEFINE_DEBUGFS_ATTRIBUTE(fops_wakelock, NULL, edgetpu_pm_debugfs_set_wakelock,
+			 "%llu\n");
+
 static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 {
 	etdev->d_entry =
@@ -1033,6 +1067,9 @@ static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 	}
 	debugfs_create_file("mappings", 0440, etdev->d_entry,
 			    etdev, &mappings_ops);
+	debugfs_create_file("syncfences", 0440, etdev->d_entry, etdev, &syncfences_ops);
+	debugfs_create_file("wakelock", 0220, etdev->d_entry, etdev,
+			    &fops_wakelock);
 #ifndef EDGETPU_FEATURE_MOBILE
 	debugfs_create_file("statusregs", 0440, etdev->d_entry, etdev,
 			    &statusregs_ops);
@@ -1230,6 +1267,10 @@ static int edgeptu_fs_add_interface(struct edgetpu_dev *etdev, struct edgetpu_de
 		return ret;
 	}
 
+	if (etiparams->name)
+		etiface->d_entry =
+			debugfs_create_symlink(etiparams->name, edgetpu_debugfs_dir,
+					       etdev->dev_name);
 	return 0;
 }
 
@@ -1265,6 +1306,7 @@ void edgetpu_fs_remove(struct edgetpu_dev *etdev)
 	for (i = 0; i < etdev->num_ifaces; i++) {
 		struct edgetpu_dev_iface *etiface = &etdev->etiface[i];
 
+		debugfs_remove(etiface->d_entry);
 		device_destroy(edgetpu_class, etiface->devno);
 		etiface->etcdev = NULL;
 		cdev_del(&etiface->cdev);
@@ -1272,20 +1314,6 @@ void edgetpu_fs_remove(struct edgetpu_dev *etdev)
 	debugfs_remove_recursive(etdev->d_entry);
 }
 
-static int syncfences_open(struct inode *inode, struct file *file)
-{
-	return single_open(file, edgetpu_sync_fence_debugfs_show,
-			   inode->i_private);
-}
-
-static const struct file_operations syncfences_ops = {
-	.open = syncfences_open,
-	.read = seq_read,
-	.llseek = seq_lseek,
-	.owner = THIS_MODULE,
-	.release = single_release,
-};
-
 static void edgetpu_debugfs_global_setup(void)
 {
 	edgetpu_debugfs_dir = debugfs_create_dir("edgetpu", NULL);
@@ -1293,9 +1321,6 @@ static void edgetpu_debugfs_global_setup(void)
 		pr_warn(DRIVER_NAME " error creating edgetpu debugfs dir\n");
 		return;
 	}
-
-	debugfs_create_file("syncfences", 0440, edgetpu_debugfs_dir, NULL,
-			    &syncfences_ops);
 }
 
 int __init edgetpu_fs_init(void)
@@ -1337,4 +1362,6 @@ struct dentry *edgetpu_fs_debugfs_dir(void)
 MODULE_DESCRIPTION("Google EdgeTPU file operations");
 MODULE_VERSION(DRIVER_VERSION);
 MODULE_LICENSE("GPL v2");
+#ifdef GIT_REPO_TAG
 MODULE_INFO(gitinfo, GIT_REPO_TAG);
+#endif
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index 58a3830..e375c0f 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -23,6 +23,13 @@
 #define EDGETPU_NUM_PREALLOCATED_DOMAINS 0
 #endif
 
+#define HAS_BEST_FIT_ALGO \
+	(IS_ENABLED(CONFIG_ANDROID) && LINUX_VERSION_CODE <= KERNEL_VERSION(5, 15, 0))
+
+#if HAS_BEST_FIT_ALGO
+#include <linux/dma-iommu.h>
+#endif
+
 struct edgetpu_iommu {
 	struct iommu_group *iommu_group;
 	/*
@@ -212,6 +219,14 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 		dev_warn(etdev->dev, "device has no iommu group\n");
 
 	etdev_warn(etdev, "AUX domains not supported\n");
+
+#if HAS_BEST_FIT_ALGO
+	/* Enable best fit algorithm to minimize fragmentation */
+	ret = iommu_dma_enable_best_fit_algo(etdev->dev);
+	if (ret)
+		etdev_warn(etdev, "Failed to enable best-fit IOVA allocator (%d)\n", ret);
+#endif
+
 	ret = check_default_domain(etdev, etiommu);
 	if (ret)
 		goto err_free;
@@ -311,18 +326,12 @@ int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
 		iommu_get_domain_for_dev(etdev->dev);
 
 	ret = get_iommu_map_params(etdev, map, context_id, &params, mmu_flags);
-
 	if (ret)
 		return ret;
 
-	if (mmu_flags & EDGETPU_MMU_64)
-		dev_warn_once(etdev->dev,
-			      "%s: 64-bit addressing is not supported",
-			      __func__);
-
 	ret = dma_map_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.nents, map->dir, map->dma_attrs);
 	if (!ret)
-		return -EINVAL;
+		return -ENOSPC;
 	map->sgt.nents = ret;
 	iova = sg_dma_address(map->sgt.sgl);
 
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 37ea27f..e8cbf9a 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -238,6 +238,7 @@ struct edgetpu_dev_iface {
 	struct edgetpu_dev *etdev; /* Pointer to core device struct */
 	dev_t devno; /* char device dev_t */
 	const char *name; /* interface specific device name */
+	struct dentry *d_entry; /* debugfs symlink if not default device name iface */
 };
 
 /* Firmware crash_type codes */
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 834318f..aea57f3 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -38,17 +38,19 @@
 /* fake-firmware could respond in a short time */
 #define KCI_TIMEOUT	(200)
 #else
-/* 5 secs. */
-#define KCI_TIMEOUT	(5000)
+/* Wait for up to 1 second for FW to respond. */
+#define KCI_TIMEOUT	(1000)
 #endif
 
-/* A macro for KCIs to leave early when the device state is known to be bad. */
-#define RETURN_ERRNO_IF_ETDEV_NOT_GOOD(kci)                                                        \
-	do {                                                                                       \
-		int ret = edgetpu_get_state_errno_locked(kci->mailbox->etdev);                     \
-		if (ret)                                                                           \
-			return ret;                                                                \
-	} while (0)
+static inline int check_etdev_state(struct edgetpu_kci *kci, char *opstring)
+{
+	int ret = edgetpu_get_state_errno_locked(kci->mailbox->etdev);
+
+	if (ret)
+		etdev_err(kci->mailbox->etdev, "%s failed: device state %u (%d)",
+			  opstring, kci->mailbox->etdev->state, ret);
+	return ret;
+}
 
 static inline u32 edgetpu_kci_queue_element_size(enum mailbox_queue_type type)
 {
@@ -201,10 +203,7 @@ static void edgetpu_reverse_kci_init(struct edgetpu_reverse_kci *rkci)
  * 2. #seq == @resp->seq:
  *   - Copy @resp, pop the head and we're done.
  * 3. #seq < @resp->seq:
- *   - Should not happen, this implies the sequence number of either entries in
- *     wait_list or responses are out-of-order, or remote didn't respond to a
- *     command. In this case, the status of response will be set to
- *     KCI_STATUS_NO_RESPONSE.
+ *   - Probable race with another context also processing KCI responses, ignore.
  *   - Pop until case 1. or 2.
  */
 static void edgetpu_kci_consume_wait_list(
@@ -225,10 +224,7 @@ static void edgetpu_kci_consume_wait_list(
 			kfree(cur);
 			break;
 		}
-		/* #seq < @resp->seq */
-		cur->resp->status = KCI_STATUS_NO_RESPONSE;
-		list_del(&cur->list);
-		kfree(cur);
+		/* #seq < @resp->seq, probable race with another consumer, let it handle. */
 	}
 
 	spin_unlock_irqrestore(&kci->wait_list_lock, flags);
@@ -246,10 +242,9 @@ edgetpu_kci_handle_response(struct edgetpu_kci *kci,
 		int ret = edgetpu_reverse_kci_add_response(kci, resp);
 
 		if (ret)
-			etdev_warn(
-				kci->mailbox->etdev,
-				"Failed to handle reverse KCI code %u (%d)\n",
-				resp->code, ret);
+			etdev_warn_ratelimited(kci->mailbox->etdev,
+					       "Failed to handle reverse KCI code %u (%d)\n",
+					       resp->code, ret);
 		return;
 	}
 	/*
@@ -642,7 +637,12 @@ int edgetpu_kci_push_cmd(struct edgetpu_kci *kci,
 
 	mutex_lock(&kci->cmd_queue_lock);
 
-	cmd->seq = kci->cur_seq;
+	/*
+	 * Only update sequence number for KCI commands.  Do not change
+	 * sequence number for responses to RKCI commands.
+	 */
+	if (!(cmd->seq & KCI_REVERSE_FLAG))
+		cmd->seq = kci->cur_seq;
 	/*
 	 * The lock ensures mailbox->cmd_queue_tail cannot be changed by
 	 * other processes (this method should be the only one to modify the
@@ -681,7 +681,8 @@ int edgetpu_kci_push_cmd(struct edgetpu_kci *kci,
 	/* triggers doorbell */
 	EDGETPU_MAILBOX_CMD_QUEUE_WRITE_SYNC(kci->mailbox, doorbell_set, 1);
 	/* bumps sequence number after the command is sent */
-	kci->cur_seq++;
+	if (!(cmd->seq & KCI_REVERSE_FLAG))
+		kci->cur_seq++;
 	ret = 0;
 out:
 	mutex_unlock(&kci->cmd_queue_lock);
@@ -707,6 +708,10 @@ static int edgetpu_kci_send_cmd_return_resp(
 	ret = edgetpu_kci_push_cmd(kci, cmd, resp);
 	if (ret)
 		return ret;
+
+	if (!resp)
+		return 0;
+
 	ret = wait_event_timeout(kci->wait_list_waitq,
 				 resp->status != KCI_STATUS_WAITING_RESPONSE,
 				 msecs_to_jiffies(KCI_TIMEOUT));
@@ -755,22 +760,11 @@ int edgetpu_kci_send_cmd(struct edgetpu_kci *kci,
 {
 	struct edgetpu_kci_response_element resp;
 
-	return edgetpu_kci_send_cmd_return_resp(kci, cmd, &resp);
-}
-
-int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
-			     u32 size, enum dma_data_direction dir)
-{
-	struct edgetpu_command_element cmd = {
-		.code = KCI_CODE_UNMAP_BUFFER,
-		.dma = {
-			.address = tpu_addr,
-			.size = size,
-			.flags = dir,
-		},
-	};
-
-	return edgetpu_kci_send_cmd(kci, &cmd);
+	/* Don't wait on a response for reverse KCI response. */
+	if (cmd->seq & KCI_REVERSE_FLAG)
+		return edgetpu_kci_send_cmd_return_resp(kci, cmd, NULL);
+	else
+		return edgetpu_kci_send_cmd_return_resp(kci, cmd, &resp);
 }
 
 int edgetpu_kci_map_log_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
@@ -810,10 +804,13 @@ int edgetpu_kci_join_group(struct edgetpu_kci *kci, u8 n_dies, u8 vid)
 		.n_dies = n_dies,
 		.vid = vid,
 	};
+	int ret;
 
 	if (!kci)
 		return -ENODEV;
-	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(kci);
+	ret = check_etdev_state(kci, "join group");
+	if (ret)
+		return ret;
 	return edgetpu_kci_send_cmd_with_data(kci, &cmd, &detail, sizeof(detail));
 }
 
@@ -822,10 +819,13 @@ int edgetpu_kci_leave_group(struct edgetpu_kci *kci)
 	struct edgetpu_command_element cmd = {
 		.code = KCI_CODE_LEAVE_GROUP,
 	};
+	int ret;
 
 	if (!kci)
 		return -ENODEV;
-	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(kci);
+	ret = check_etdev_state(kci, "leave group");
+	if (ret)
+		return ret;
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
@@ -926,6 +926,9 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 
 fw_unlock:
 	edgetpu_firmware_unlock(etdev);
+
+	if (ret)
+		etdev_warn_once(etdev, "get firmware usage stats failed: %d", ret);
 	return ret;
 }
 
@@ -933,10 +936,11 @@ int edgetpu_kci_update_usage_locked(struct edgetpu_dev *etdev)
 {
 #define EDGETPU_USAGE_BUFFER_SIZE	4096
 	struct edgetpu_command_element cmd = {
-		.code = KCI_CODE_GET_USAGE,
+		.code = KCI_CODE_GET_USAGE_V2,
 		.dma = {
 			.address = 0,
 			.size = 0,
+			.flags = EDGETPU_USAGE_METRIC_VERSION,
 		},
 	};
 	struct edgetpu_coherent_mem mem;
@@ -952,13 +956,22 @@ int edgetpu_kci_update_usage_locked(struct edgetpu_dev *etdev)
 		return ret;
 	}
 
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+retry_v1:
+	if (etdev->usage_stats && etdev->usage_stats->use_metrics_v1)
+		cmd.code = KCI_CODE_GET_USAGE_V1;
 	cmd.dma.address = mem.tpu_addr;
 	cmd.dma.size = EDGETPU_USAGE_BUFFER_SIZE;
 	memset(mem.vaddr, 0, sizeof(struct edgetpu_usage_header));
 	ret = edgetpu_kci_send_cmd_return_resp(etdev->kci, &cmd, &resp);
 
-	if (ret == KCI_ERROR_UNIMPLEMENTED || ret == KCI_ERROR_UNAVAILABLE)
+	if (ret == KCI_ERROR_UNIMPLEMENTED || ret == KCI_ERROR_UNAVAILABLE) {
+		if (etdev->usage_stats && !etdev->usage_stats->use_metrics_v1) {
+			etdev->usage_stats->use_metrics_v1 = true;
+			goto retry_v1;
+		}
 		etdev_dbg(etdev, "firmware does not report usage\n");
+	}
 	else if (ret == KCI_ERROR_OK)
 		edgetpu_usage_stats_process_buffer(etdev, mem.vaddr);
 	else if (ret != -ETIMEDOUT)
@@ -1024,10 +1037,11 @@ int edgetpu_kci_get_debug_dump(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
-int edgetpu_kci_open_device(struct edgetpu_kci *kci, u32 mailbox_map, s16 vcid, bool first_open)
+int edgetpu_kci_open_device(struct edgetpu_kci *kci, u32 mailbox_map, u32 client_priv, s16 vcid,
+			    bool first_open)
 {
 	const struct edgetpu_kci_open_device_detail detail = {
-		.mailbox_map = mailbox_map,
+		.client_priv = client_priv,
 		.vcid = vcid,
 		.flags = (mailbox_map << 1) | first_open,
 	};
@@ -1037,10 +1051,13 @@ int edgetpu_kci_open_device(struct edgetpu_kci *kci, u32 mailbox_map, s16 vcid,
 			.flags = mailbox_map,
 		},
 	};
+	int ret;
 
 	if (!kci)
 		return -ENODEV;
-	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(kci);
+	ret = check_etdev_state(kci, "open device");
+	if (ret)
+		return ret;
 	if (vcid < 0)
 		return edgetpu_kci_send_cmd(kci, &cmd);
 	return edgetpu_kci_send_cmd_with_data(kci, &cmd, &detail, sizeof(detail));
@@ -1054,10 +1071,13 @@ int edgetpu_kci_close_device(struct edgetpu_kci *kci, u32 mailbox_map)
 			.flags = mailbox_map,
 		},
 	};
+	int ret;
 
 	if (!kci)
 		return -ENODEV;
-	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(kci);
+	ret = check_etdev_state(kci, "close device");
+	if (ret)
+		return ret;
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
@@ -1091,3 +1111,36 @@ int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block)
 
 	return edgetpu_kci_send_cmd(etdev->kci, &cmd);
 }
+
+int edgetpu_kci_firmware_tracing_level(struct edgetpu_dev *etdev, unsigned long level,
+				       unsigned long *active_level)
+{
+	struct edgetpu_command_element cmd = {
+		.code = KCI_CODE_FIRMWARE_TRACING_LEVEL,
+		.dma = {
+			.flags = (u32)level,
+		},
+	};
+	struct edgetpu_kci_response_element resp;
+	int ret;
+
+	ret = edgetpu_kci_send_cmd_return_resp(etdev->kci, &cmd, &resp);
+	if (ret == KCI_ERROR_OK)
+		*active_level = resp.retval;
+
+	return ret;
+}
+
+int edgetpu_kci_resp_rkci_ack(struct edgetpu_dev *etdev,
+			      struct edgetpu_kci_response_element *rkci_cmd)
+{
+	struct edgetpu_command_element cmd = {
+		.seq = rkci_cmd->seq,
+		.code = KCI_CODE_RKCI_ACK,
+	};
+
+	if (!etdev->kci)
+		return -ENODEV;
+
+	return edgetpu_kci_send_cmd(etdev->kci, &cmd);
+}
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index cc8bff3..2a2ff16 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -60,7 +60,11 @@ struct edgetpu_dma_descriptor {
 };
 
 struct edgetpu_command_element {
-	u64 seq;	 /* set by edgetpu_kci_push_cmd() */
+	/*
+	 * Set by edgetpu_kci_push_cmd() in case of KCI cmd and copied from
+	 * the RKCI cmd in case of RKCI response.
+	 */
+	u64 seq;
 	u16 code;
 	u16 reserved[3]; /* explicit padding, does not affect alignment */
 	struct edgetpu_dma_descriptor dma;
@@ -72,7 +76,8 @@ struct edgetpu_kci_response_element {
 	/*
 	 * Reserved for host use - firmware can't touch this.
 	 * If a value is written here it will be discarded and overwritten
-	 * during response processing.
+	 * during response processing. However, when repurposed as an RKCI
+	 * command, the FW can set this field.
 	 */
 	u16 status;
 	/*
@@ -107,9 +112,16 @@ enum edgetpu_kci_code {
 	KCI_CODE_OPEN_DEVICE = 9,
 	KCI_CODE_CLOSE_DEVICE = 10,
 	KCI_CODE_FIRMWARE_INFO = 11,
-	KCI_CODE_GET_USAGE = 12,
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+	KCI_CODE_GET_USAGE_V1 = 12,
 	KCI_CODE_NOTIFY_THROTTLING = 13,
 	KCI_CODE_BLOCK_BUS_SPEED_CONTROL = 14,
+	/* 15..18 not implemented in this branch */
+	KCI_CODE_FIRMWARE_TRACING_LEVEL = 19,
+	/* 20 not implemented in this branch */
+	KCI_CODE_GET_USAGE_V2 = 21,
+
+	KCI_CODE_RKCI_ACK = 256,
 };
 
 /*
@@ -207,8 +219,8 @@ struct edgetpu_kci_device_group_detail {
 };
 
 struct edgetpu_kci_open_device_detail {
-	/* The bit map of mailboxes to be opened. */
-	u16 mailbox_map;
+	/* The client privilege level. */
+	u16 client_priv;
 	/*
 	 * Virtual context ID @mailbox_id is associated to.
 	 * For device groups with @mailbox_detachable attribute the mailbox attached to the group
@@ -284,14 +296,6 @@ int edgetpu_kci_send_cmd(struct edgetpu_kci *kci,
 			 struct edgetpu_command_element *cmd);
 
 /*
- * Sends the "Unmap Buffer Sync" command and waits for remote response.
- *
- * Returns the code of response, or a negative errno on error.
- */
-int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
-			     u32 size, enum dma_data_direction dir);
-
-/*
  * Sends a FIRMWARE_INFO command and expects a response with a
  * edgetpu_fw_info struct filled out, including what firmware type is running,
  * along with build CL and time.
@@ -377,7 +381,8 @@ int edgetpu_kci_get_debug_dump(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
  * You usually shouldn't call this directly - consider using
  * edgetpu_mailbox_activate() or edgetpu_mailbox_activate_bulk() instead.
  */
-int edgetpu_kci_open_device(struct edgetpu_kci *kci, u32 mailbox_map, s16 vcid, bool first_open);
+int edgetpu_kci_open_device(struct edgetpu_kci *kci, u32 mailbox_map, u32 client_priv, s16 vcid,
+			    bool first_open);
 
 /*
  * Inform the firmware that the VII mailboxes included in @mailbox_map are closed.
@@ -405,4 +410,18 @@ int edgetpu_kci_notify_throttling(struct edgetpu_dev *etdev, u32 level);
  */
 int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block);
 
+/* Set the firmware tracing level. */
+int edgetpu_kci_firmware_tracing_level(struct edgetpu_dev *etdev, unsigned long level,
+				       unsigned long *active_level);
+
+/*
+ * Send an ack to the FW after handling a reverse KCI request.
+ *
+ * The FW may wait for a response from the kernel for an RKCI request so a
+ * response could be sent as an ack.
+ */
+int edgetpu_kci_resp_rkci_ack(struct edgetpu_dev *etdev,
+			      struct edgetpu_kci_response_element *rkci_cmd);
+
+
 #endif /* __EDGETPU_KCI_H__ */
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 5fac2a2..cd9a6b2 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -1102,7 +1102,8 @@ void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *g
 	edgetpu_mailbox_external_free(group);
 }
 
-static int edgetpu_mailbox_external_enable_by_id(struct edgetpu_client *client, int mailbox_id)
+static int edgetpu_mailbox_external_enable_by_id(struct edgetpu_client *client, int mailbox_id,
+						 u32 client_priv)
 {
 	int ret;
 
@@ -1115,7 +1116,7 @@ static int edgetpu_mailbox_external_enable_by_id(struct edgetpu_client *client,
 
 	etdev_dbg(client->etdev, "Enabling mailbox: %d\n", mailbox_id);
 
-	ret = edgetpu_mailbox_activate(client->etdev, mailbox_id, -1, false);
+	ret = edgetpu_mailbox_activate(client->etdev, mailbox_id, client_priv, -1, false);
 	if (ret)
 		etdev_err(client->etdev, "Activate mailbox %d failed: %d", mailbox_id, ret);
 	else
@@ -1162,7 +1163,8 @@ int edgetpu_mailbox_activate_external_mailbox(struct edgetpu_device_group *group
 	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);
+	ret = edgetpu_mailbox_activate_bulk(ext_mailbox->etdev, mbox_map,
+					    group->mbox_attr.client_priv, vcid, false);
 
 	if (ret)
 		etdev_err(group->etdev, "Activate mailbox bulk failed: %d", ret);
@@ -1198,12 +1200,13 @@ void edgetpu_mailbox_deactivate_external_mailbox(struct edgetpu_device_group *gr
 }
 
 int edgetpu_mailbox_enable_ext(struct edgetpu_client *client, int mailbox_id,
-			       struct edgetpu_external_mailbox_req *ext_mailbox_req)
+			       struct edgetpu_external_mailbox_req *ext_mailbox_req,
+			       u32 client_priv)
 {
 	if (mailbox_id == EDGETPU_MAILBOX_ID_USE_ASSOC)
 		return edgetpu_mailbox_external_alloc_enable(client, ext_mailbox_req);
 	else
-		return edgetpu_mailbox_external_enable_by_id(client, mailbox_id);
+		return edgetpu_mailbox_external_enable_by_id(client, mailbox_id, client_priv);
 }
 
 int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id)
@@ -1214,16 +1217,16 @@ int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id)
 		return edgetpu_mailbox_external_disable_by_id(client, mailbox_id);
 }
 
-int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, s16 vcid,
-				  bool first_open)
+int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, u32 client_priv,
+				  s16 vcid, bool first_open)
 {
 	struct edgetpu_handshake *eh = &etdev->mailbox_manager->open_devices;
 	int ret = 0;
 
 	mutex_lock(&eh->lock);
 	if (mailbox_map & ~eh->fw_state)
-		ret = edgetpu_kci_open_device(etdev->kci, mailbox_map & ~eh->fw_state, vcid,
-					      first_open);
+		ret = edgetpu_kci_open_device(etdev->kci, mailbox_map & ~eh->fw_state, client_priv,
+					      vcid, first_open);
 	if (!ret) {
 		eh->state |= mailbox_map;
 		eh->fw_state |= mailbox_map;
@@ -1240,9 +1243,10 @@ int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, s1
 
 }
 
-int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, s16 vcid, bool first_open)
+int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, u32 client_priv, s16 vcid,
+			     bool first_open)
 {
-	return edgetpu_mailbox_activate_bulk(etdev, BIT(mailbox_id), vcid, first_open);
+	return edgetpu_mailbox_activate_bulk(etdev, BIT(mailbox_id), client_priv, vcid, first_open);
 }
 
 void edgetpu_mailbox_deactivate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map)
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index 115341d..b03ae24 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -349,7 +349,8 @@ int edgetpu_mailbox_p2p_batch(struct edgetpu_mailbox_manager *mgr, uint n,
  * 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);
+			       struct edgetpu_external_mailbox_req *ext_mailbox_req,
+			       u32 client_priv);
 
 /*
  * Notify firmware of an external mailboxes becoming inactive.
@@ -362,8 +363,8 @@ int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id);
  * Returns what edgetpu_kci_open_device() returned.
  * Caller ensures device is powered on.
  */
-int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, s16 vcid,
-				  bool first_open);
+int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, u32 client_priv,
+				  s16 vcid, bool first_open);
 
 /*
  * Activates @mailbox_id, OPEN_DEVICE KCI will be sent.
@@ -374,7 +375,8 @@ int edgetpu_mailbox_activate_bulk(struct edgetpu_dev *etdev, u32 mailbox_map, s1
  * Returns what edgetpu_kci_open_device() returned.
  * Caller ensures device is powered on.
  */
-int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, s16 vcid, bool first_open);
+int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, u32 client_priv, s16 vcid,
+			     bool first_open);
 
 /*
  * Similar to edgetpu_mailbox_activate_bulk() but sends CLOSE_DEVICE KCI with the @mailbox_map
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index af6bcb7..979c109 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -22,6 +22,8 @@
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
 
+static struct edgetpu_dev *edgetpu_debug_pointer;
+
 /*
  * Log and trace buffers at the beginning of the remapped region,
  * pool memory afterwards.
@@ -157,7 +159,7 @@ int edgetpu_chip_acquire_ext_mailbox(struct edgetpu_client *client,
 		mutex_unlock(&etmdev->tz_mailbox_lock);
 		return -EBUSY;
 	}
-	ret = edgetpu_mailbox_enable_ext(client, EDGETPU_TZ_MAILBOX_ID, NULL);
+	ret = edgetpu_mailbox_enable_ext(client, EDGETPU_TZ_MAILBOX_ID, NULL, 0);
 	if (!ret)
 		etmdev->secure_client = client;
 	mutex_unlock(&etmdev->tz_mailbox_lock);
@@ -220,6 +222,20 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
+static void edgetpu_platform_parse_pmu(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct device *dev = etdev->dev;
+	u32 reg;
+
+	if (of_find_property(dev->of_node, "pmu-status-base", NULL) &&
+	    !of_property_read_u32_index(dev->of_node, "pmu-status-base", 0, &reg)) {
+		etmdev->pmu_status = devm_ioremap(dev, reg, 0x4);
+		if (!etmdev->pmu_status)
+			etdev_info(etdev, "Using ACPM for blk status query\n");
+	}
+}
+
 static int edgetpu_platform_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
 {
 	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
@@ -311,6 +327,17 @@ edgetpu_mobile_platform_set_fw_ctx_memory(struct edgetpu_mobile_platform_dev *et
 	return 0;
 }
 
+static inline const char *get_driver_commit(void)
+{
+#if IS_ENABLED(CONFIG_MODULE_SCMVERSION)
+	return THIS_MODULE->scmversion ?: "scmversion missing";
+#elif defined(GIT_REPO_TAG)
+	return GIT_REPO_TAG;
+#else
+	return "Unknown";
+#endif
+}
+
 static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 					 struct edgetpu_mobile_platform_dev *etmdev)
 {
@@ -397,6 +424,8 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 	if (ret)
 		dev_warn(dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
 
+	edgetpu_platform_parse_pmu(etmdev);
+
 	etmdev->log_mem = devm_kcalloc(dev, etdev->num_cores, sizeof(*etmdev->log_mem), GFP_KERNEL);
 	if (!etmdev->log_mem) {
 		ret = -ENOMEM;
@@ -434,10 +463,13 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 		}
 	}
 
-	dev_info(dev, "%s edgetpu initialized. Build: %s", etdev->dev_name, GIT_REPO_TAG);
+	dev_info(dev, "%s edgetpu initialized. Build: %s", etdev->dev_name, get_driver_commit());
+
 	/* Turn the device off unless a client request is already received. */
 	edgetpu_pm_shutdown(etdev, false);
 
+	edgetpu_debug_pointer = etdev;
+
 	return 0;
 out_destroy_fw:
 	edgetpu_mobile_firmware_destroy(etdev);
@@ -473,6 +505,9 @@ static int edgetpu_mobile_platform_remove(struct platform_device *pdev)
 	edgetpu_platform_cleanup_fw_region(etmdev);
 	edgetpu_pm_put(etdev->pm);
 	edgetpu_pm_shutdown(etdev, true);
-	mobile_pm_destroy(etdev);
+	edgetpu_mobile_pm_destroy(etdev);
+
+	edgetpu_debug_pointer = NULL;
+
 	return 0;
 }
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index 9d41571..335eb94 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -45,8 +45,8 @@ struct edgetpu_mobile_platform_pwr {
 	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);
+	/* Block shutdown status callback, may be NULL */
+	bool (*is_block_down)(struct edgetpu_dev *etdev);
 
 	/* Firmware shutdown callback. Must be implemented */
 	void (*firmware_down)(struct edgetpu_dev *etdev);
@@ -105,6 +105,8 @@ struct edgetpu_mobile_platform_dev {
 	int n_irq;
 	/* Array of IRQ numbers */
 	int *irq;
+	/* PMU status base address for block status, maybe NULL */
+	void __iomem *pmu_status;
 
 	/* callbacks for chip-dependent implementations */
 
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index a71232d..40d41ff 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -53,9 +53,13 @@ static int edgetpu_pm_get_locked(struct edgetpu_pm *etpm)
 	int ret = 0;
 
 	if (!power_up_count) {
-		ret = etpm->p->handlers->power_up(etpm);
-		if (!ret)
-			edgetpu_mailbox_restore_active_mailbox_queues(etpm->etdev);
+		if (etpm->p->power_down_pending) {
+			etpm->p->power_down_pending = false;
+		} else {
+			ret = etpm->p->handlers->power_up(etpm);
+			if (!ret)
+				edgetpu_mailbox_restore_active_mailbox_queues(etpm->etdev);
+		}
 	}
 	if (ret)
 		etpm->p->power_up_count--;
@@ -103,7 +107,6 @@ int edgetpu_pm_get(struct edgetpu_pm *etpm)
 		return 0;
 
 	mutex_lock(&etpm->p->lock);
-	etpm->p->power_down_pending = false;
 	ret = edgetpu_pm_get_locked(etpm);
 	mutex_unlock(&etpm->p->lock);
 
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index dbd283f..0c163e0 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -56,17 +56,6 @@ int edgetpu_thermal_suspend(struct device *dev);
 int edgetpu_thermal_resume(struct device *dev);
 
 /*
- * 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.
  */
@@ -77,15 +66,4 @@ static inline bool edgetpu_thermal_is_suspended(struct edgetpu_thermal *thermal)
 	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 ba93d49..9934ca6 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -74,6 +74,7 @@ int edgetpu_usage_add(struct edgetpu_dev *etdev, struct tpu_usage *tpu_usage)
 	if (!ustats)
 		return 0;
 
+	/* Note: as of metrics v2 the cluster_id is always zero and is ignored. */
 	etdev_dbg(etdev, "%s: uid=%u state=%u dur=%u", __func__,
 		  tpu_usage->uid, tpu_usage->power_state,
 		  tpu_usage->duration_us);
@@ -125,63 +126,78 @@ static void edgetpu_utilization_update(
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
-static void edgetpu_counter_update(
-	struct edgetpu_dev *etdev,
-	struct edgetpu_usage_counter *counter)
+static void edgetpu_counter_update(struct edgetpu_dev *etdev, struct edgetpu_usage_counter *counter,
+				   uint version)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	uint component = version > 1 ? counter->component_id : 0;
 
 	if (!ustats)
 		return;
 
-	etdev_dbg(etdev, "%s: type=%d value=%llu\n", __func__,
-		  counter->type, counter->value);
+	etdev_dbg(etdev, "%s: type=%d value=%llu comp=%u\n", __func__, counter->type,
+		  counter->value, component);
 
 	mutex_lock(&ustats->usage_stats_lock);
 	if (counter->type >= 0 && counter->type < EDGETPU_COUNTER_COUNT)
-		ustats->counter[counter->type] += counter->value;
+		ustats->counter[counter->type][component] += counter->value;
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
-static void edgetpu_counter_clear(
-	struct edgetpu_dev *etdev,
-	enum edgetpu_usage_counter_type counter_type)
+static void edgetpu_counter_clear(struct edgetpu_dev *etdev,
+				  enum edgetpu_usage_counter_type counter_type)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int i;
 
-	if (!ustats)
-		return;
 	if (counter_type >= EDGETPU_COUNTER_COUNT)
 		return;
 
 	mutex_lock(&ustats->usage_stats_lock);
-	ustats->counter[counter_type] = 0;
+	for (i = 0; i < EDGETPU_TPU_CLUSTER_COUNT; i++)
+		ustats->counter[counter_type][i] = 0;
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
-static void edgetpu_max_watermark_update(
-	struct edgetpu_dev *etdev,
-	struct edgetpu_usage_max_watermark *max_watermark)
+static void edgetpu_max_watermark_update(struct edgetpu_dev *etdev,
+					 struct edgetpu_usage_max_watermark *max_watermark,
+					 uint version)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	uint component = version > 1 ? max_watermark->component_id : 0;
 
 	if (!ustats)
 		return;
 
-	etdev_dbg(etdev, "%s: type=%d value=%llu\n", __func__,
-		  max_watermark->type, max_watermark->value);
+	etdev_dbg(etdev, "%s: type=%d value=%llu comp=%u\n", __func__, max_watermark->type,
+		  max_watermark->value, component);
 
 	if (max_watermark->type < 0 ||
 	    max_watermark->type >= EDGETPU_MAX_WATERMARK_TYPE_COUNT)
 		return;
 
 	mutex_lock(&ustats->usage_stats_lock);
-	if (max_watermark->value > ustats->max_watermark[max_watermark->type])
-		ustats->max_watermark[max_watermark->type] =
+	if (max_watermark->value > ustats->max_watermark[max_watermark->type][component])
+		ustats->max_watermark[max_watermark->type][component] =
 			max_watermark->value;
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
+static void edgetpu_max_watermark_clear(struct edgetpu_dev *etdev,
+					enum edgetpu_usage_max_watermark_type max_watermark_type)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int i;
+
+	if (max_watermark_type < 0 || max_watermark_type >= EDGETPU_MAX_WATERMARK_TYPE_COUNT)
+		return;
+
+	mutex_lock(&ustats->usage_stats_lock);
+	for (i = 0; i < EDGETPU_TPU_CLUSTER_COUNT; i++)
+		ustats->max_watermark[max_watermark_type][i] = 0;
+	mutex_unlock(&ustats->usage_stats_lock);
+}
+
 static void edgetpu_thread_stats_update(
 	struct edgetpu_dev *etdev,
 	struct edgetpu_thread_stats *thread_stats)
@@ -241,20 +257,44 @@ out:
 
 void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 {
-	struct edgetpu_usage_header *header = buf;
-	struct edgetpu_usage_metric *metric =
-		(struct edgetpu_usage_metric *)(header + 1);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	struct edgetpu_usage_metric *metric;
+	uint metric_size;
+	uint num_metrics;
+	uint version;
 	int i;
 
-	etdev_dbg(etdev, "%s: n=%u sz=%u", __func__,
-		  header->num_metrics, header->metric_size);
-	if (header->metric_size != sizeof(struct edgetpu_usage_metric)) {
-		etdev_dbg(etdev, "%s: expected sz=%zu, discard", __func__,
-			  sizeof(struct edgetpu_usage_metric));
+	if (!ustats)
+		return;
+
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+	if (ustats->use_metrics_v1) {
+		struct edgetpu_usage_header_v1 *header = buf;
+
+		metric_size = header->metric_size;
+		num_metrics = header->num_metrics;
+		version = 1;
+		metric = (struct edgetpu_usage_metric *)(header + 1);
+	} else {
+		struct edgetpu_usage_header *header = buf;
+
+		metric_size = header->metric_size;
+		num_metrics = header->num_metrics;
+		version = header->version;
+		metric = (struct edgetpu_usage_metric *)((char *)header + header->header_bytes);
+	}
+
+	etdev_dbg(etdev, "%s: v=%u n=%u sz=%u", __func__, version, num_metrics, metric_size);
+	if (metric_size < EDGETPU_USAGE_METRIC_SIZE_V1) {
+		etdev_warn_once(etdev, "fw metric size %u less than minimum %u",
+				metric_size, EDGETPU_USAGE_METRIC_SIZE_V1);
 		return;
 	}
 
-	for (i = 0; i < header->num_metrics; i++) {
+	if (metric_size > sizeof(struct edgetpu_usage_metric))
+		etdev_dbg(etdev, "fw metrics are later version with unknown fields");
+
+	for (i = 0; i < num_metrics; i++) {
 		switch (metric->type) {
 		case EDGETPU_METRIC_TYPE_TPU_USAGE:
 			edgetpu_usage_add(etdev, &metric->tpu_usage);
@@ -264,19 +304,16 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 				etdev, &metric->component_activity);
 			break;
 		case EDGETPU_METRIC_TYPE_COUNTER:
-			edgetpu_counter_update(etdev, &metric->counter);
+			edgetpu_counter_update(etdev, &metric->counter, version);
 			break;
 		case EDGETPU_METRIC_TYPE_MAX_WATERMARK:
-			edgetpu_max_watermark_update(
-				etdev, &metric->max_watermark);
+			edgetpu_max_watermark_update(etdev, &metric->max_watermark, version);
 			break;
 		case EDGETPU_METRIC_TYPE_THREAD_STATS:
-			edgetpu_thread_stats_update(
-				etdev, &metric->thread_stats);
+			edgetpu_thread_stats_update(etdev, &metric->thread_stats);
 			break;
 		case EDGETPU_METRIC_TYPE_DVFS_FREQUENCY_INFO:
-			edgetpu_dvfs_frequency_update(
-				etdev, metric->dvfs_frequency_info);
+			edgetpu_dvfs_frequency_update(etdev, metric->dvfs_frequency_info);
 			break;
 		default:
 			etdev_dbg(etdev, "%s: %d: skip unknown type=%u",
@@ -284,7 +321,7 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 			break;
 		}
 
-		metric++;
+		metric = (struct edgetpu_usage_metric *)((char *)metric + metric_size);
 	}
 }
 
@@ -304,36 +341,72 @@ int edgetpu_usage_get_utilization(struct edgetpu_dev *etdev,
 	return val;
 }
 
-static int64_t edgetpu_usage_get_counter(
-	struct edgetpu_dev *etdev,
-	enum edgetpu_usage_counter_type counter_type)
+/*
+ * Resyncs firmware stats and formats the requested counter in the supplied buffer.
+ *
+ * If @report_per_cluster is true, and if the firmware implements metrics V2 or higher,
+ * then one value is formatted per cluster (for chips with only one cluster only one value is
+ * formatted).
+ *
+ * Returns the number of bytes written to buf.
+ */
+static ssize_t edgetpu_usage_format_counter(struct edgetpu_dev *etdev, char *buf,
+					    enum edgetpu_usage_counter_type counter_type,
+					    bool report_per_cluster)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-	int64_t val;
+	uint ncomponents = report_per_cluster && !etdev->usage_stats->use_metrics_v1 ?
+		EDGETPU_TPU_CLUSTER_COUNT : 1;
+	uint i;
+	ssize_t ret = 0;
 
 	if (counter_type >= EDGETPU_COUNTER_COUNT)
-		return -1;
+		return 0;
 	edgetpu_kci_update_usage(etdev);
 	mutex_lock(&ustats->usage_stats_lock);
-	val = ustats->counter[counter_type];
+	for (i = 0; i < ncomponents; i++) {
+		if (i)
+			ret += scnprintf(buf + ret, PAGE_SIZE - ret, " ");
+		ret += scnprintf(buf + ret, PAGE_SIZE - ret, "%llu",
+				 ustats->counter[counter_type][i]);
+	}
 	mutex_unlock(&ustats->usage_stats_lock);
-	return val;
+	ret += scnprintf(buf + ret, PAGE_SIZE - ret, "\n");
+	return ret;
 }
 
-static int64_t edgetpu_usage_get_max_watermark(
-	struct edgetpu_dev *etdev,
-	enum edgetpu_usage_max_watermark_type max_watermark_type)
+/*
+ * Resyncs firmware stats and formats the requested max watermark in the supplied buffer.
+ *
+ * If @report_per_cluster is true, and if the firmware implements metrics V2 or higher,
+ * then one value is formatted per cluster (for chips with only one cluster only one value is
+ * formatted).
+ *
+ * Returns the number of bytes written to buf.
+ */
+static ssize_t edgetpu_usage_format_max_watermark(
+	struct edgetpu_dev *etdev, char *buf,
+	enum edgetpu_usage_max_watermark_type max_watermark_type, bool report_per_cluster)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-	int64_t val;
+	uint ncomponents = report_per_cluster && !etdev->usage_stats->use_metrics_v1 ?
+		EDGETPU_TPU_CLUSTER_COUNT : 1;
+	uint i;
+	ssize_t ret = 0;
 
 	if (max_watermark_type >= EDGETPU_MAX_WATERMARK_TYPE_COUNT)
-		return -1;
+		return 0;
 	edgetpu_kci_update_usage(etdev);
 	mutex_lock(&ustats->usage_stats_lock);
-	val = ustats->max_watermark[max_watermark_type];
+	for (i = 0; i < ncomponents; i++) {
+		if (i)
+			ret += scnprintf(buf + ret, PAGE_SIZE - ret, " ");
+		ret += scnprintf(buf + ret, PAGE_SIZE - ret, "%llu",
+				 ustats->max_watermark[max_watermark_type][i]);
+	}
 	mutex_unlock(&ustats->usage_stats_lock);
-	return val;
+	ret += scnprintf(buf + ret, PAGE_SIZE - ret, "\n");
+	return ret;
 }
 
 static ssize_t tpu_usage_show(struct device *dev,
@@ -447,11 +520,8 @@ static ssize_t tpu_active_cycle_count_show(struct device *dev,
 					   char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_TPU_ACTIVE_CYCLES);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_TPU_ACTIVE_CYCLES, false);
 }
 
 static ssize_t tpu_active_cycle_count_store(struct device *dev,
@@ -472,11 +542,8 @@ static ssize_t tpu_throttle_stall_count_show(struct device *dev,
 					     char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_TPU_THROTTLE_STALLS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_TPU_THROTTLE_STALLS, false);
 }
 
 static ssize_t tpu_throttle_stall_count_store(struct device *dev,
@@ -497,11 +564,8 @@ static ssize_t inference_count_show(struct device *dev,
 				    struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_INFERENCES);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_INFERENCES, true);
 }
 
 static ssize_t inference_count_store(struct device *dev,
@@ -517,21 +581,15 @@ static ssize_t inference_count_store(struct device *dev,
 static DEVICE_ATTR(inference_count, 0664, inference_count_show,
 		   inference_count_store);
 
-static ssize_t tpu_op_count_show(struct device *dev,
-				 struct device_attribute *attr, char *buf)
+static ssize_t tpu_op_count_show(struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_TPU_OPS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_TPU_OPS, true);
 }
 
-static ssize_t tpu_op_count_store(struct device *dev,
-				  struct device_attribute *attr,
-				  const char *buf,
-				  size_t count)
+static ssize_t tpu_op_count_store(struct device *dev, struct device_attribute *attr,
+				  const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 
@@ -540,22 +598,16 @@ static ssize_t tpu_op_count_store(struct device *dev,
 }
 static DEVICE_ATTR(tpu_op_count, 0664, tpu_op_count_show, tpu_op_count_store);
 
-static ssize_t param_cache_hit_count_show(struct device *dev,
-					  struct device_attribute *attr,
+static ssize_t param_cache_hit_count_show(struct device *dev, struct device_attribute *attr,
 					  char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_PARAM_CACHE_HITS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_PARAM_CACHE_HITS, false);
 }
 
-static ssize_t param_cache_hit_count_store(struct device *dev,
-					   struct device_attribute *attr,
-					   const char *buf,
-					   size_t count)
+static ssize_t param_cache_hit_count_store(struct device *dev, struct device_attribute *attr,
+					   const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 
@@ -565,22 +617,16 @@ static ssize_t param_cache_hit_count_store(struct device *dev,
 static DEVICE_ATTR(param_cache_hit_count, 0664, param_cache_hit_count_show,
 		   param_cache_hit_count_store);
 
-static ssize_t param_cache_miss_count_show(struct device *dev,
-					   struct device_attribute *attr,
+static ssize_t param_cache_miss_count_show(struct device *dev, struct device_attribute *attr,
 					   char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_PARAM_CACHE_MISSES);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_PARAM_CACHE_MISSES, false);
 }
 
-static ssize_t param_cache_miss_count_store(struct device *dev,
-					    struct device_attribute *attr,
-					    const char *buf,
-					    size_t count)
+static ssize_t param_cache_miss_count_store(struct device *dev, struct device_attribute *attr,
+					    const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 
@@ -590,22 +636,16 @@ static ssize_t param_cache_miss_count_store(struct device *dev,
 static DEVICE_ATTR(param_cache_miss_count, 0664, param_cache_miss_count_show,
 		   param_cache_miss_count_store);
 
-static ssize_t context_preempt_count_show(struct device *dev,
-					  struct device_attribute *attr,
+static ssize_t context_preempt_count_show(struct device *dev, struct device_attribute *attr,
 					  char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev,
-					EDGETPU_COUNTER_CONTEXT_PREEMPTS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_CONTEXT_PREEMPTS, true);
 }
 
-static ssize_t context_preempt_count_store(struct device *dev,
-					   struct device_attribute *attr,
-					   const char *buf,
-					   size_t count)
+static ssize_t context_preempt_count_store(struct device *dev, struct device_attribute *attr,
+					   const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 
@@ -619,10 +659,8 @@ static ssize_t hardware_preempt_count_show(struct device *dev, struct device_att
 					   char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev, EDGETPU_COUNTER_HARDWARE_PREEMPTS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_HARDWARE_PREEMPTS, true);
 }
 
 static ssize_t hardware_preempt_count_store(struct device *dev, struct device_attribute *attr,
@@ -640,10 +678,9 @@ static ssize_t hardware_ctx_save_time_show(struct device *dev, struct device_att
 					   char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev, EDGETPU_COUNTER_HARDWARE_CTX_SAVE_TIME_US);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_HARDWARE_CTX_SAVE_TIME_US,
+					    true);
 }
 
 static ssize_t hardware_ctx_save_time_store(struct device *dev, struct device_attribute *attr,
@@ -661,10 +698,9 @@ static ssize_t scalar_fence_wait_time_show(struct device *dev, struct device_att
 					   char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev, EDGETPU_COUNTER_SCALAR_FENCE_WAIT_TIME_US);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_SCALAR_FENCE_WAIT_TIME_US,
+					    true);
 }
 
 static ssize_t scalar_fence_wait_time_store(struct device *dev, struct device_attribute *attr,
@@ -679,13 +715,11 @@ static DEVICE_ATTR(scalar_fence_wait_time, 0664, scalar_fence_wait_time_show,
 		   scalar_fence_wait_time_store);
 
 static ssize_t long_suspend_count_show(struct device *dev, struct device_attribute *attr,
-					   char *buf)
+				       char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_counter(etdev, EDGETPU_COUNTER_LONG_SUSPEND);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_LONG_SUSPEND, false);
 }
 
 static ssize_t long_suspend_count_store(struct device *dev, struct device_attribute *attr,
@@ -699,15 +733,53 @@ static ssize_t long_suspend_count_store(struct device *dev, struct device_attrib
 static DEVICE_ATTR(long_suspend_count, 0664, long_suspend_count_show,
 		   long_suspend_count_store);
 
+#if EDGETPU_TPU_CLUSTER_COUNT > 1
+static ssize_t reconfigurations_show(struct device *dev, struct device_attribute *attr, char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_RECONFIGURATIONS, false);
+}
+
+static ssize_t reconfigurations_store(struct device *dev, struct device_attribute *attr,
+				      const char *buf, size_t count)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	edgetpu_counter_clear(etdev, EDGETPU_COUNTER_RECONFIGURATIONS);
+	return count;
+}
+static DEVICE_ATTR(reconfigurations, 0664, reconfigurations_show, reconfigurations_store);
+
+static ssize_t preempt_reconfigurations_show(struct device *dev, struct device_attribute *attr,
+					     char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	return edgetpu_usage_format_counter(etdev, buf, EDGETPU_COUNTER_PREEMPT_RECONFIGURATIONS,
+					    false);
+}
+
+static ssize_t preempt_reconfigurations_store(struct device *dev, struct device_attribute *attr,
+					      const char *buf, size_t count)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	edgetpu_counter_clear(etdev, EDGETPU_COUNTER_PREEMPT_RECONFIGURATIONS);
+	return count;
+}
+static DEVICE_ATTR(preempt_reconfigurations, 0664, preempt_reconfigurations_show,
+		   preempt_reconfigurations_store);
+#endif /* EDGETPU_TPU_CLUSTER_COUNT > 1 */
+
+
 static ssize_t outstanding_commands_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_max_watermark(
-			etdev, EDGETPU_MAX_WATERMARK_OUT_CMDS);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_max_watermark(etdev, buf, EDGETPU_MAX_WATERMARK_OUT_CMDS,
+						  false);
 }
 
 static ssize_t outstanding_commands_max_store(
@@ -715,14 +787,8 @@ static ssize_t outstanding_commands_max_store(
 	const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-
-	if (ustats) {
-		mutex_lock(&ustats->usage_stats_lock);
-		ustats->max_watermark[EDGETPU_MAX_WATERMARK_OUT_CMDS] = 0;
-		mutex_unlock(&ustats->usage_stats_lock);
-	}
 
+	edgetpu_max_watermark_clear(etdev, EDGETPU_MAX_WATERMARK_OUT_CMDS);
 	return count;
 }
 static DEVICE_ATTR(outstanding_commands_max, 0664,
@@ -733,11 +799,9 @@ static ssize_t preempt_depth_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_max_watermark(
-			etdev, EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_max_watermark(etdev, buf, EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH,
+						  true);
 }
 
 static ssize_t preempt_depth_max_store(
@@ -745,14 +809,8 @@ static ssize_t preempt_depth_max_store(
 	const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-
-	if (ustats) {
-		mutex_lock(&ustats->usage_stats_lock);
-		ustats->max_watermark[EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH] = 0;
-		mutex_unlock(&ustats->usage_stats_lock);
-	}
 
+	edgetpu_max_watermark_clear(etdev, EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH);
 	return count;
 }
 static DEVICE_ATTR(preempt_depth_max, 0664, preempt_depth_max_show,
@@ -762,11 +820,10 @@ static ssize_t hardware_ctx_save_time_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_max_watermark(
-			etdev, EDGETPU_MAX_WATERMARK_HARDWARE_CTX_SAVE_TIME_US);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_max_watermark(etdev, buf,
+						  EDGETPU_MAX_WATERMARK_HARDWARE_CTX_SAVE_TIME_US,
+						  true);
 }
 
 static ssize_t hardware_ctx_save_time_max_store(
@@ -774,14 +831,8 @@ static ssize_t hardware_ctx_save_time_max_store(
 	const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-
-	if (ustats) {
-		mutex_lock(&ustats->usage_stats_lock);
-		ustats->max_watermark[EDGETPU_MAX_WATERMARK_HARDWARE_CTX_SAVE_TIME_US] = 0;
-		mutex_unlock(&ustats->usage_stats_lock);
-	}
 
+	edgetpu_max_watermark_clear(etdev, EDGETPU_MAX_WATERMARK_HARDWARE_CTX_SAVE_TIME_US);
 	return count;
 }
 static DEVICE_ATTR(hardware_ctx_save_time_max, 0664, hardware_ctx_save_time_max_show,
@@ -791,11 +842,9 @@ static ssize_t scalar_fence_wait_time_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_max_watermark(
-			etdev, EDGETPU_MAX_WATERMARK_SCALAR_FENCE_WAIT_TIME_US);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_max_watermark(
+			etdev, buf, EDGETPU_MAX_WATERMARK_SCALAR_FENCE_WAIT_TIME_US, true);
 }
 
 static ssize_t scalar_fence_wait_time_max_store(
@@ -803,14 +852,8 @@ static ssize_t scalar_fence_wait_time_max_store(
 	const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-
-	if (ustats) {
-		mutex_lock(&ustats->usage_stats_lock);
-		ustats->max_watermark[EDGETPU_MAX_WATERMARK_SCALAR_FENCE_WAIT_TIME_US] = 0;
-		mutex_unlock(&ustats->usage_stats_lock);
-	}
 
+	edgetpu_max_watermark_clear(etdev, EDGETPU_MAX_WATERMARK_SCALAR_FENCE_WAIT_TIME_US);
 	return count;
 }
 static DEVICE_ATTR(scalar_fence_wait_time_max, 0664, scalar_fence_wait_time_max_show,
@@ -820,11 +863,9 @@ static ssize_t suspend_time_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	int64_t val;
 
-	val = edgetpu_usage_get_max_watermark(
-			etdev, EDGETPU_MAX_WATERMARK_SUSPEND_TIME_US);
-	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+	return edgetpu_usage_format_max_watermark(etdev, buf, EDGETPU_MAX_WATERMARK_SUSPEND_TIME_US,
+						  false);
 }
 
 static ssize_t suspend_time_max_store(
@@ -832,14 +873,8 @@ static ssize_t suspend_time_max_store(
 	const char *buf, size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-
-	if (ustats) {
-		mutex_lock(&ustats->usage_stats_lock);
-		ustats->max_watermark[EDGETPU_MAX_WATERMARK_SUSPEND_TIME_US] = 0;
-		mutex_unlock(&ustats->usage_stats_lock);
-	}
 
+	edgetpu_max_watermark_clear(etdev, EDGETPU_MAX_WATERMARK_SUSPEND_TIME_US);
 	return count;
 }
 static DEVICE_ATTR(suspend_time_max, 0664, suspend_time_max_show,
@@ -900,6 +935,10 @@ static struct attribute *usage_stats_dev_attrs[] = {
 	&dev_attr_hardware_ctx_save_time.attr,
 	&dev_attr_scalar_fence_wait_time.attr,
 	&dev_attr_long_suspend_count.attr,
+#if EDGETPU_TPU_CLUSTER_COUNT > 1
+	&dev_attr_reconfigurations.attr,
+	&dev_attr_preempt_reconfigurations.attr,
+#endif
 	&dev_attr_outstanding_commands_max.attr,
 	&dev_attr_preempt_depth_max.attr,
 	&dev_attr_hardware_ctx_save_time_max.attr,
@@ -912,6 +951,7 @@ static struct attribute *usage_stats_dev_attrs[] = {
 static const struct attribute_group usage_stats_attr_group = {
 	.attrs = usage_stats_dev_attrs,
 };
+
 void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_usage_stats *ustats;
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index a60b107..2d97043 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -10,9 +10,31 @@
 #include <linux/hashtable.h>
 #include <linux/mutex.h>
 
+/* The highest version of usage metrics handled by this driver. */
+#define EDGETPU_USAGE_METRIC_VERSION	2
+
+/* Max # of TPU clusters accounted for in the highest supported metrics version. */
+#define EDGETPU_USAGE_CLUSTERS_MAX	3
+
+/*
+ * Size in bytes of usage metric v1.
+ * If fewer bytes than this are received then discard the invalid buffer.
+ * This size also identifies the fw response as v1; subsequent versions will add another field
+ * with the version number.
+ */
+#define EDGETPU_USAGE_METRIC_SIZE_V1	20
+
+/* v1 metric header struct. */
+struct edgetpu_usage_header_v1 {
+	uint32_t num_metrics;		/* Number of metrics being reported */
+	uint32_t metric_size;		/* Size of each metric struct */
+};
+
 /* Header struct in the metric buffer. */
 /* Must be kept in sync with firmware struct UsageTrackerHeader */
 struct edgetpu_usage_header {
+	uint16_t header_bytes;		/* Number of bytes in this header */
+	uint16_t version;		/* Metrics version */
 	uint32_t num_metrics;		/* Number of metrics being reported */
 	uint32_t metric_size;		/* Size of each metric struct */
 };
@@ -20,15 +42,25 @@ struct edgetpu_usage_header {
 /*
  * Encapsulate TPU core usage information of a specific application for a
  * specific power state.
- * Must be kept in sync with firmware struct TpuUsage.
+ * Must be kept in sync with firmware struct CoreUsage.
  */
 struct tpu_usage {
 	/* Unique identifier of the application. */
 	int32_t uid;
 	/* The power state of the device (values are chip dependent) */
+	/* Now called operating_point in FW. */
 	uint32_t power_state;
 	/* Duration of usage in microseconds. */
 	uint32_t duration_us;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Compute Core: TPU cluster ID. */
+	/* Called core_id in FW. */
+	/* Note: as of metrics v2 the cluster_id is always zero and is ignored. */
+	uint8_t cluster_id;
+	/* Reserved.  Filling out the next 32-bit boundary. */
+	uint8_t reserved[3];
 };
 
 /*
@@ -38,9 +70,13 @@ struct tpu_usage {
 enum edgetpu_usage_component {
 	/* The device as a whole */
 	EDGETPU_USAGE_COMPONENT_DEVICE = 0,
-	/* Just the TPU core */
+	/* Just the TPU core (scalar core and tiles) */
 	EDGETPU_USAGE_COMPONENT_TPU = 1,
-	EDGETPU_USAGE_COMPONENT_COUNT = 2, /* number of components above */
+	/* Control core (ARM Cortex-R52 CPU) */
+	/* Note: this component is not reported as of metrics v2. */
+	EDGETPU_USAGE_COMPONENT_CONTROLCORE = 2,
+
+	EDGETPU_USAGE_COMPONENT_COUNT = 3, /* number of components above */
 };
 
 /*
@@ -62,7 +98,7 @@ enum edgetpu_usage_counter_type {
 	EDGETPU_COUNTER_TPU_ACTIVE_CYCLES = 0,
 	/* Number of stalls caused by throttling. */
 	EDGETPU_COUNTER_TPU_THROTTLE_STALLS = 1,
-	/* Number of graph invocations. */
+	/* Number of graph invocations. (Now called kWorkload in FW.) */
 	EDGETPU_COUNTER_INFERENCES = 2,
 	/* Number of TPU offload op invocations. */
 	EDGETPU_COUNTER_TPU_OPS = 3,
@@ -81,7 +117,18 @@ enum edgetpu_usage_counter_type {
 	/* Number of times (firmware)suspend function takes longer than SLA time. */
 	EDGETPU_COUNTER_LONG_SUSPEND = 10,
 
-	EDGETPU_COUNTER_COUNT = 11, /* number of counters above */
+	/* The following counters are added in metrics v2. */
+
+	/* Counter 11 not used on TPU. */
+	EDGETPU_COUNTER_CONTEXT_SWITCHES = 11,
+
+	/* Number of TPU Cluster Reconfigurations. */
+	EDGETPU_COUNTER_RECONFIGURATIONS = 12,
+
+	/* Number of TPU Cluster Reconfigurations motivated exclusively by a preemption. */
+	EDGETPU_COUNTER_PREEMPT_RECONFIGURATIONS = 13,
+
+	EDGETPU_COUNTER_COUNT = 14, /* number of counters above */
 };
 
 /* Generic counter. Only reported if it has a value larger than 0. */
@@ -91,6 +138,11 @@ struct __packed edgetpu_usage_counter {
 
 	/* Accumulated value since last initialization. */
 	uint64_t value;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Reporting component. */
+	uint8_t component_id;
 };
 
 /* Defines different max watermarks we track. */
@@ -121,15 +173,22 @@ struct __packed edgetpu_usage_max_watermark {
 	 * non-mobile, firmware boot on mobile).
 	 */
 	uint64_t value;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Reporting component. */
+	uint8_t component_id;
 };
 
 /* An enum to identify the tracked firmware threads. */
 /* Must be kept in sync with firmware enum class UsageTrackerThreadId. */
 enum edgetpu_usage_threadid {
-	/* Individual thread IDs are not tracked. */
+	/* Individual thread IDs do not have identifiers assigned. */
+
+	/* Thread ID 14 is not used for TPU */
 
 	/* Number of task identifiers. */
-	EDGETPU_FW_THREAD_COUNT = 12,
+	EDGETPU_FW_THREAD_COUNT = 17,
 };
 
 /* Statistics related to a single thread in firmware. */
@@ -173,11 +232,13 @@ struct edgetpu_usage_metric {
 #define UID_HASH_BITS 3
 
 struct edgetpu_usage_stats {
+	/* if true the current firmware only implements metrics V1 */
+	bool use_metrics_v1;
 	DECLARE_HASHTABLE(uid_hash_table, UID_HASH_BITS);
 	/* component utilization values reported by firmware */
 	int32_t component_utilization[EDGETPU_USAGE_COMPONENT_COUNT];
-	int64_t counter[EDGETPU_COUNTER_COUNT];
-	int64_t max_watermark[EDGETPU_MAX_WATERMARK_TYPE_COUNT];
+	int64_t counter[EDGETPU_COUNTER_COUNT][EDGETPU_USAGE_CLUSTERS_MAX];
+	int64_t max_watermark[EDGETPU_MAX_WATERMARK_TYPE_COUNT][EDGETPU_USAGE_CLUSTERS_MAX];
 	int32_t thread_stack_max[EDGETPU_FW_THREAD_COUNT];
 	struct mutex usage_stats_lock;
 };
diff --git a/drivers/edgetpu/edgetpu.h b/drivers/edgetpu/edgetpu.h
index ad9a2ce..2be97fe 100644
--- a/drivers/edgetpu/edgetpu.h
+++ b/drivers/edgetpu/edgetpu.h
@@ -182,6 +182,7 @@ struct edgetpu_mailbox_attr {
 	__u32 cmdq_tail_doorbell: 1; /* auto doorbell on cmd queue tail move */
 	/* Type of memory partitions to be used for this group, exact meaning is chip-dependent. */
 	__u32 partition_type    : 1;
+	__u32 client_priv : 1; /* client privilege level */
 };
 
 /*
diff --git a/drivers/edgetpu/mobile-firmware.c b/drivers/edgetpu/mobile-firmware.c
index c1d58d6..42f982f 100644
--- a/drivers/edgetpu/mobile-firmware.c
+++ b/drivers/edgetpu/mobile-firmware.c
@@ -281,35 +281,6 @@ static int mobile_firmware_gsa_authenticate(struct edgetpu_mobile_platform_dev *
 	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 void mobile_firmware_setup_ssmt(struct edgetpu_dev *etdev)
 {
 	int i;
@@ -331,8 +302,6 @@ static void mobile_firmware_setup_ssmt(struct edgetpu_dev *etdev)
 	 * Reset the table to zeroes if running non-secure firmware, since the SSMT
 	 * will be in clamped mode and we want all memory accesses to go to the
 	 * default page table.
-	 *
-	 * TODO(b/204384254) Confirm SSMT setup for Rio
 	 */
 	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
 		int val;
@@ -355,9 +324,6 @@ static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	/* 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);
-
 	mobile_firmware_setup_ssmt(etdev);
 
 	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
@@ -390,6 +356,7 @@ static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	phys_addr_t image_start, image_end, carveout_start, carveout_end;
 	bool image_config_changed;
+	struct mobile_image_header *hdr;
 
 	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
 		etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
@@ -404,6 +371,11 @@ static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 		return -ENOMEM;
 	}
 
+	hdr = (struct mobile_image_header *)fw_buf->vaddr;
+	if (hdr->Magic != EDGETPU_MOBILE_FW_MAGIC)
+		etdev_warn(etdev, "Invalid firmware header magic value %#08x\n",
+			   hdr->Magic);
+
 	/* fetch the firmware versions */
 	image_config = fw_buf->vaddr + MOBILE_IMAGE_CONFIG_OFFSET;
 	memcpy(&etdev->fw_version, &image_config->firmware_versions,
diff --git a/drivers/edgetpu/mobile-firmware.h b/drivers/edgetpu/mobile-firmware.h
index 27ee048..05102c9 100644
--- a/drivers/edgetpu/mobile-firmware.h
+++ b/drivers/edgetpu/mobile-firmware.h
@@ -75,6 +75,9 @@ struct mobile_image_header {
 	struct mobile_image_config ImageConfig;
 };
 
+/* Value of Magic field above: 'TPUF' as a 32-bit LE int */
+#define EDGETPU_MOBILE_FW_MAGIC	0x46555054
+
 int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev);
 void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev);
 
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 52a43ab..40db11c 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -21,6 +21,7 @@
 #include "edgetpu-mailbox.h"
 #include "edgetpu-mobile-platform.h"
 #include "edgetpu-pm.h"
+#include "edgetpu-thermal.h"
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
 
@@ -40,6 +41,10 @@ module_param(power_state, int, 0660);
 
 #define MAX_VOLTAGE_VAL 1250000
 
+#define BLOCK_DOWN_RETRY_TIMES 50
+#define BLOCK_DOWN_MIN_DELAY_US 1000
+#define BLOCK_DOWN_MAX_DELAY_US 1500
+
 enum edgetpu_pwr_state edgetpu_active_states[EDGETPU_NUM_STATES] = {
 	TPU_ACTIVE_UUD,
 	TPU_ACTIVE_SUD,
@@ -64,6 +69,7 @@ static int mobile_pwr_state_init(struct device *dev)
 		ret = pm_runtime_get_sync(dev);
 		if (ret) {
 			pm_runtime_put_noidle(dev);
+			pm_runtime_disable(dev);
 			dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
 			return ret;
 		}
@@ -74,6 +80,7 @@ static int mobile_pwr_state_init(struct device *dev)
 		dev_err(dev, "error initializing tpu state: %d\n", ret);
 		if (curr_state > TPU_OFF)
 			pm_runtime_put_sync(dev);
+		pm_runtime_disable(dev);
 		return ret;
 	}
 
@@ -255,8 +262,6 @@ static int mobile_pwr_state_set_locked(struct edgetpu_mobile_platform_dev *etmde
 			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;
@@ -432,7 +437,27 @@ 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));
+	int ret;
+
+	if (platform_pwr->is_block_down) {
+		int times = 0;
+
+		do {
+			if (platform_pwr->is_block_down(etdev))
+				break;
+			usleep_range(BLOCK_DOWN_MIN_DELAY_US, BLOCK_DOWN_MAX_DELAY_US);
+		} while (++times < BLOCK_DOWN_RETRY_TIMES);
+		if (times >= BLOCK_DOWN_RETRY_TIMES && !platform_pwr->is_block_down(etdev))
+			return -EAGAIN;
+	}
+
+	if (edgetpu_thermal_is_suspended(etdev->thermal)) {
+		etdev_warn_ratelimited(etdev,
+				       "power up rejected due to device thermal limit exceeded");
+		return -EAGAIN;
+	}
+
+	ret = mobile_pwr_state_set(etpm->etdev, mobile_get_initial_pwr_state(etdev->dev));
 
 	etdev_info(etpm->etdev, "Powering up\n");
 
@@ -600,6 +625,7 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 	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;
+	int curr_state;
 
 	ret = mobile_pwr_state_init(dev);
 	if (ret)
@@ -641,7 +667,19 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 
 	if (platform_pwr->after_create)
 		ret = platform_pwr->after_create(etdev);
+	if (ret)
+		goto err_debugfs_remove;
+
+	return 0;
 
+err_debugfs_remove:
+	debugfs_remove_recursive(platform_pwr->debugfs_dir);
+	/* pm_runtime_{enable,get_sync} were called in mobile_pwr_state_init */
+
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	if (curr_state > TPU_OFF)
+		pm_runtime_put_sync(dev);
+	pm_runtime_disable(dev);
 	return ret;
 }
 
@@ -668,17 +706,17 @@ static struct edgetpu_pm_handlers mobile_pm_handlers = {
 	.power_down = mobile_power_down,
 };
 
-int mobile_pm_create(struct edgetpu_dev *etdev)
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev)
 {
 	return edgetpu_pm_create(etdev, &mobile_pm_handlers);
 }
 
-void mobile_pm_destroy(struct edgetpu_dev *etdev)
+void edgetpu_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)
+void edgetpu_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;
@@ -726,7 +764,6 @@ static void mobile_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 		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;
 	}
@@ -741,7 +778,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, u16 bts_val)
+void edgetpu_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 4299b49..b0f8c3d 100644
--- a/drivers/edgetpu/mobile-pm.h
+++ b/drivers/edgetpu/mobile-pm.h
@@ -47,6 +47,8 @@ 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,
+	/* The above codes have been deprecated. */
+
 	RKCI_CODE_PM_QOS_BTS = RKCI_CHIP_CODE_FIRST + 3,
 };
 
@@ -68,7 +70,7 @@ enum mobile_reverse_kci_code {
  * chipsets.
  * Needs to be called after the devices's platform_pwr struct has been initialized.
  */
-int mobile_pm_create(struct edgetpu_dev *etdev);
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev);
 
 /*
  * Wrapper for chip-specific implementation.
@@ -79,12 +81,12 @@ 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);
+void edgetpu_mobile_pm_destroy(struct edgetpu_dev *etdev);
 
 /* Set required QoS value for the edgetpu device. */
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
+void edgetpu_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, u16 bts_val);
+void edgetpu_mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val);
 
 #endif /* __MOBILE_PM_H__ */