Merge branch 'whitechapel' into android-gs-pixel-5.10

* whitechapel:
  edgetpu: detachable VII queues
  edgetpu: guard mbox queue mmap with wakelock
  edgetpu: set wakelock to null on cleanup
  edgetpu: APIs that bind KCI and removable mailbox
  edgetpu: abrolhos: cap thermal state change requests
  edgetpu: adopt new wakelock interface
  edgetpu: add chip files for wakelock
  edgetpu: add edgetpu-wakelock.c
  edgetpu: abrolhos: do SMC on power up/down

Signed-off-by: Nrithya Kanakasabapathy <nrithya@google.com>
Change-Id: I2f2d8927e3edd2c0d804d7984221c2af102a6a37

14 files changed, 520 insertions, 192 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index f6c7d96..f1e7f56 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -12,7 +12,7 @@ endif
 
 edgetpu-fw-objs := edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-shared-fw.o
 edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o $(edgetpu-fw-objs)
-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 $(edgetpu-objs)
+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 $(edgetpu-objs)
 CFLAGS_abrolhos-fs.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-core.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device.o := -DCONFIG_ABROLHOS=1
@@ -24,3 +24,4 @@ CFLAGS_abrolhos-pm.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-thermal.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-debug-dump.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-usage-stats.o := -DCONFIG_ABROLHOS=1
+CFLAGS_abrolhos-wakelock.o := -DCONFIG_ABROLHOS=1
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index 6af7f73..b30f95f 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -23,7 +23,7 @@ abrolhos-objs	:= abrolhos-core.o abrolhos-debug-dump.o \
 		   abrolhos-device-group.o abrolhos-device.o \
 		   abrolhos-firmware.o abrolhos-fs.o abrolhos-iommu.o \
 		   abrolhos-platform.o abrolhos-pm.o abrolhos-thermal.o \
-		   abrolhos-usage-stats.o \
+		   abrolhos-usage-stats.o abrolhos-wakelock.o \
 		   $(edgetpu-objs)
 
 KBUILD_OPTIONS += CONFIG_ABROLHOS=m
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index b71e45d..61f9bd8 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -10,6 +10,7 @@
 #include <linux/gsa/gsa_tpu.h>
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
+#include <linux/soc/samsung/exynos-smc.h>
 
 #include "abrolhos-platform.h"
 #include "abrolhos-pm.h"
@@ -26,6 +27,8 @@
 
 #include "edgetpu-pm.c"
 
+#define TPU_SMC_ID			(0x15)
+
 /*
  * Encode INT/MIF values as a 16 bit pair in the 32-bit return value
  * (in units of MHz, to provide enough range)
@@ -94,6 +97,12 @@ static int abrolhos_pwr_state_set(void *data, u64 val)
 			dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
 			return ret;
 		}
+		ret = exynos_smc(SMC_PROTECTION_SET, 0, TPU_SMC_ID,
+				 SMC_PROTECTION_ENABLE);
+		if (ret)
+			dev_warn(dev,
+				 "exynos_smc protection enable returned %d\n",
+				 ret);
 	}
 
 	ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, (unsigned long)val);
@@ -104,6 +113,13 @@ static int abrolhos_pwr_state_set(void *data, u64 val)
 	}
 
 	if (curr_state != TPU_OFF && val == TPU_OFF) {
+		ret = exynos_smc(SMC_PROTECTION_SET, 0, TPU_SMC_ID,
+				 SMC_PROTECTION_DISABLE);
+		if (ret)
+			dev_warn(dev,
+				 "exynos_smc protection disable returned %d\n",
+				 ret);
+
 		ret = pm_runtime_put_sync(dev);
 		if (ret) {
 			dev_err(dev, "%s: pm_runtime_put_sync returned %d\n",
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index 2740d08..bfa28cf 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -82,6 +82,18 @@ static int edgetpu_set_cur_state(struct thermal_cooling_device *cdev,
 	mutex_lock(&cooling->lock);
 	pwr_state = state_mapping[state_original];
 	if (state_original != cooling->cooling_state) {
+		/*
+		 * Cap the minimum state we request here.
+		 * We cannot go to states below SUD until firmware/runtime
+		 * handshake is added.
+		 */
+		if (pwr_state < TPU_ACTIVE_SUD) {
+			dev_warn_ratelimited(
+				dev, "Unable to go to state %lu, going to %d",
+				pwr_state, TPU_ACTIVE_SUD);
+			pwr_state = TPU_ACTIVE_SUD;
+		}
+
 		ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, pwr_state);
 		if (ret) {
 			dev_err(dev, "error setting tpu policy: %d\n", ret);
diff --git a/drivers/edgetpu/abrolhos-wakelock.c b/drivers/edgetpu/abrolhos-wakelock.c
new file mode 100644
index 0000000..ab61fb2
--- /dev/null
+++ b/drivers/edgetpu/abrolhos-wakelock.c
@@ -0,0 +1,2 @@
+// SPDX-License-Identifier: GPL-2.0
+#include "edgetpu-wakelock.c"
diff --git a/drivers/edgetpu/abrolhos/config.h b/drivers/edgetpu/abrolhos/config.h
index 262b3a5..fa52c01 100644
--- a/drivers/edgetpu/abrolhos/config.h
+++ b/drivers/edgetpu/abrolhos/config.h
@@ -14,6 +14,8 @@
 
 #define EDGETPU_HAS_MULTI_GROUPS
 
+#define EDGETPU_HAS_WAKELOCK
+
 /*
  * A remapped data region is available. This will be accessible by the R52
  * regardless of active context and is typically used for logging buffer and
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index e1ac96f..e4ae403 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -33,12 +33,9 @@
 #include "edgetpu-mmu.h"
 #include "edgetpu-telemetry.h"
 #include "edgetpu-usage-stats.h"
+#include "edgetpu-wakelock.h"
 #include "edgetpu.h"
 
-#define UNLOCK(client) mutex_unlock(&client->group_lock)
-#define LOCK_IN_GROUP(client)                           \
-	({ mutex_lock(&client->group_lock); client->group ? 0 : -EINVAL; })
-
 static atomic_t single_dev_count = ATOMIC_INIT(-1);
 
 static int edgetpu_mmap_compat(struct edgetpu_client *client,
@@ -64,39 +61,42 @@ static int edgetpu_mmap_compat(struct edgetpu_client *client,
 	return ret;
 }
 
-static void edgetpu_vma_open(struct vm_area_struct *vma)
+/*
+ * Returns the wakelock event by mmap offset. Returns EDGETPU_WAKELOCK_EVENT_END
+ * if the offset does not correspond to a wakelock event.
+ */
+static enum edgetpu_wakelock_event mmap_wakelock_event(unsigned long pgoff)
 {
-	struct edgetpu_client *client = vma->vm_private_data;
-
-	switch (vma->vm_pgoff) {
+	switch (pgoff) {
 	case 0:
+		return EDGETPU_WAKELOCK_EVENT_FULL_CSR;
 	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
-		mutex_lock(&client->wakelock.lock);
-		client->wakelock.csr_map_count++;
-		mutex_unlock(&client->wakelock.lock);
-		break;
+		return EDGETPU_WAKELOCK_EVENT_MBOX_CSR;
+	case EDGETPU_MMAP_CMD_QUEUE_OFFSET >> PAGE_SHIFT:
+		return EDGETPU_WAKELOCK_EVENT_CMD_QUEUE;
+	case EDGETPU_MMAP_RESP_QUEUE_OFFSET >> PAGE_SHIFT:
+		return EDGETPU_WAKELOCK_EVENT_RESP_QUEUE;
+	default:
+		return EDGETPU_WAKELOCK_EVENT_END;
 	}
 }
 
+static void edgetpu_vma_open(struct vm_area_struct *vma)
+{
+	struct edgetpu_client *client = vma->vm_private_data;
+	enum edgetpu_wakelock_event evt = mmap_wakelock_event(vma->vm_pgoff);
+
+	if (evt != EDGETPU_WAKELOCK_EVENT_END)
+		edgetpu_wakelock_inc_event(client->wakelock, evt);
+}
+
 static void edgetpu_vma_close(struct vm_area_struct *vma)
 {
 	struct edgetpu_client *client = vma->vm_private_data;
+	enum edgetpu_wakelock_event evt = mmap_wakelock_event(vma->vm_pgoff);
 
-	switch (vma->vm_pgoff) {
-	case 0:
-	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
-		mutex_lock(&client->wakelock.lock);
-		if (!client->wakelock.csr_map_count)
-			etdev_warn(client->etdev,
-				   "unbalanced vma_close on CSR mapping\n");
-		else
-			client->wakelock.csr_map_count--;
-		etdev_dbg(client->etdev,
-			  "%s: unmap CSRS. pgoff = %lX count = %u\n", __func__,
-			  vma->vm_pgoff, client->wakelock.csr_map_count);
-		mutex_unlock(&client->wakelock.lock);
-		break;
-	}
+	if (evt != EDGETPU_WAKELOCK_EVENT_END)
+		edgetpu_wakelock_dec_event(client->wakelock, evt);
 }
 
 static const struct vm_operations_struct edgetpu_vma_ops = {
@@ -104,11 +104,11 @@ static const struct vm_operations_struct edgetpu_vma_ops = {
 	.close = edgetpu_vma_close,
 };
 
-
 /* Map exported device CSRs or queue into user space. */
 int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 {
 	int ret = 0;
+	enum edgetpu_wakelock_event evt;
 
 	if (vma->vm_start & ~PAGE_MASK) {
 		etdev_dbg(client->etdev,
@@ -128,14 +128,15 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 
 	/* If backward compat map all CSRs */
 	if (!vma->vm_pgoff) {
-		mutex_lock(&client->wakelock.lock);
-		if (!client->wakelock.req_count)
-			ret = -EAGAIN;
-		else
+		evt = EDGETPU_WAKELOCK_EVENT_FULL_CSR;
+		if (edgetpu_wakelock_inc_event(client->wakelock, evt)) {
 			ret = edgetpu_mmap_compat(client, vma);
-		if (!ret)
-			client->wakelock.csr_map_count++;
-		mutex_unlock(&client->wakelock.lock);
+			if (ret)
+				edgetpu_wakelock_dec_event(client->wakelock,
+							   evt);
+		} else {
+			ret = -EAGAIN;
+		}
 		return ret;
 	}
 
@@ -147,41 +148,33 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 		return edgetpu_mmap_telemetry_buffer(
 			client->etdev, EDGETPU_TELEMETRY_TRACE, vma);
 
+	evt = mmap_wakelock_event(vma->vm_pgoff);
+	if (evt == EDGETPU_WAKELOCK_EVENT_END)
+		return -EINVAL;
+	if (!edgetpu_wakelock_inc_event(client->wakelock, evt))
+		return -EAGAIN;
+
+	mutex_lock(&client->group_lock);
+	if (!client->group) {
+		ret = -EINVAL;
+		goto out_unlock;
+	}
 	switch (vma->vm_pgoff) {
 	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
-		mutex_lock(&client->wakelock.lock);
-		if (!client->wakelock.req_count) {
-			ret = -EAGAIN;
-		} else {
-			ret = LOCK_IN_GROUP(client);
-			if (!ret)
-				ret = edgetpu_mmap_csr(client->group, vma);
-			UNLOCK(client);
-		}
-		if (!ret)
-			client->wakelock.csr_map_count++;
-		etdev_dbg(client->etdev, "%s: mmap CSRS. count = %u ret = %d\n",
-			  __func__, client->wakelock.csr_map_count, ret);
-		mutex_unlock(&client->wakelock.lock);
+		ret = edgetpu_mmap_csr(client->group, vma);
 		break;
 	case EDGETPU_MMAP_CMD_QUEUE_OFFSET >> PAGE_SHIFT:
-		ret = LOCK_IN_GROUP(client);
-		if (!ret)
-			ret = edgetpu_mmap_queue(client->group,
-						 MAILBOX_CMD_QUEUE, vma);
-		UNLOCK(client);
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma);
 		break;
 	case EDGETPU_MMAP_RESP_QUEUE_OFFSET >> PAGE_SHIFT:
-		ret = LOCK_IN_GROUP(client);
-		if (!ret)
-			ret = edgetpu_mmap_queue(client->group,
-						 MAILBOX_RESP_QUEUE, vma);
-		UNLOCK(client);
-		break;
-	default:
-		ret = -EINVAL;
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE,
+					 vma);
 		break;
 	}
+out_unlock:
+	mutex_unlock(&client->group_lock);
+	if (ret)
+		edgetpu_wakelock_dec_event(client->wakelock, evt);
 	return ret;
 }
 
@@ -324,6 +317,11 @@ struct edgetpu_client *edgetpu_client_add(struct edgetpu_dev *etdev)
 	client = kzalloc(sizeof(*client), GFP_KERNEL);
 	if (!client)
 		return ERR_PTR(-ENOMEM);
+	client->wakelock = edgetpu_wakelock_alloc(etdev);
+	if (!client->wakelock) {
+		kfree(client);
+		return ERR_PTR(-ENOMEM);
+	}
 
 	/* Allow entire CSR space to be mmap()'ed using 1.0 interface */
 	client->reg_window.start_reg_offset = 0;
@@ -332,9 +330,6 @@ struct edgetpu_client *edgetpu_client_add(struct edgetpu_dev *etdev)
 	client->tgid = current->tgid;
 	client->etdev = etdev;
 	mutex_init(&client->group_lock);
-	mutex_init(&client->wakelock.lock);
-	/* Initialize client wakelock state to "acquired" */
-	client->wakelock.req_count = 1;
 	/* equivalent to edgetpu_client_get() */
 	refcount_set(&client->count, 1);
 	return client;
@@ -368,6 +363,13 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	 */
 	if (client->group)
 		edgetpu_device_group_leave(client);
+	edgetpu_wakelock_free(client->wakelock);
+	/*
+	 * It should be impossible to access client->wakelock after this cleanup
+	 * procedure. Set to NULL to cause kernel panic if use-after-free does
+	 * happen.
+	 */
+	client->wakelock = NULL;
 	edgetpu_client_put(client);
 }
 
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 5035314..59bc79c 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -94,9 +94,13 @@ static int edgetpu_kci_leave_group_worker(struct kci_worker_param *param)
 
 static int edgetpu_group_kci_open_device(struct edgetpu_device_group *group)
 {
-	u8 mailbox_id = edgetpu_group_context_id_locked(group);
-	int ret = edgetpu_kci_open_device(group->etdev->kci, BIT(mailbox_id));
+	u8 mailbox_id;
+	int ret;
 
+	if (edgetpu_group_mailbox_detached_locked(group))
+		return 0;
+	mailbox_id = edgetpu_group_context_id_locked(group);
+	ret = edgetpu_kci_open_device(group->etdev->kci, BIT(mailbox_id));
 	/*
 	 * This should only happen when the FW hasn't driven this KCI, log once
 	 * to prevent log storm.
@@ -436,21 +440,32 @@ static int edgetpu_dev_add_group(struct edgetpu_dev *etdev,
 				 struct edgetpu_device_group *group)
 {
 	struct edgetpu_list_group *l = kmalloc(sizeof(*l), GFP_KERNEL);
+	int ret;
 
 	if (!l)
 		return -ENOMEM;
 	mutex_lock(&etdev->groups_lock);
 	if (etdev->group_join_lockout) {
-		mutex_unlock(&etdev->groups_lock);
-		kfree(l);
-		return -EAGAIN;
+		ret = -EAGAIN;
+		goto error_unlock;
 	}
+#ifndef EDGETPU_HAS_MULTI_GROUPS
+	if (etdev->n_groups >= 1) {
+		ret = -EBUSY;
+		goto error_unlock;
+	}
+#endif /* !EDGETPU_HAS_MULTI_GROUPS */
 	l->grp = edgetpu_device_group_get(group);
 	list_add_tail(&l->list, &etdev->groups);
 	etdev->n_groups++;
 
 	mutex_unlock(&etdev->groups_lock);
 	return 0;
+
+error_unlock:
+	mutex_unlock(&etdev->groups_lock);
+	kfree(l);
+	return ret;
 }
 
 void edgetpu_device_group_put(struct edgetpu_device_group *group)
@@ -1511,66 +1526,58 @@ void edgetpu_fatal_error_notify(struct edgetpu_dev *etdev)
 	mutex_unlock(&etdev->groups_lock);
 }
 
-void edgetpu_group_detach_mailbox(struct edgetpu_device_group *group)
+void edgetpu_group_detach_mailbox_locked(struct edgetpu_device_group *group)
 {
-	struct edgetpu_mailbox_manager *mgr = group->etdev->mailbox_manager;
-	struct edgetpu_mailbox *mailbox;
-
 	if (!group->mailbox_detachable)
 		return;
-	mutex_lock(&group->lock);
-	if (edgetpu_group_mailbox_detached_locked(group)) {
-		mutex_unlock(&group->lock);
+	if (edgetpu_group_mailbox_detached_locked(group))
 		return;
-	}
-	if (edgetpu_device_group_is_finalized(group))
-		edgetpu_group_kci_close_device(group);
-	mailbox = group->vii.mailbox;
-	group->vii.mailbox = NULL;
-	edgetpu_device_group_put(mailbox->internal.group);
-	edgetpu_mailbox_remove(mgr, mailbox);
+	edgetpu_mailbox_remove_vii(&group->vii);
 	edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
 	if (group->etdomain->token != EDGETPU_DOMAIN_TOKEN_END)
 		group->context_id =
 			EDGETPU_CONTEXT_DOMAIN_TOKEN | group->etdomain->token;
 	else
 		group->context_id = EDGETPU_CONTEXT_INVALID;
+}
 
+void edgetpu_group_close_and_detach_mailbox(struct edgetpu_device_group *group)
+{
+	mutex_lock(&group->lock);
+	if (edgetpu_device_group_is_finalized(group))
+		edgetpu_group_kci_close_device(group);
+	edgetpu_group_detach_mailbox_locked(group);
 	mutex_unlock(&group->lock);
 }
 
-int edgetpu_group_attach_mailbox(struct edgetpu_device_group *group)
+int edgetpu_group_attach_mailbox_locked(struct edgetpu_device_group *group)
 {
-	struct edgetpu_mailbox_manager *mgr = group->etdev->mailbox_manager;
-	struct edgetpu_mailbox *mailbox;
-	uint ctx_id;
-	int ret = 0;
+	int ret;
 
 	if (!group->mailbox_detachable)
 		return 0;
-	mutex_lock(&group->lock);
 	if (!edgetpu_group_mailbox_detached_locked(group))
-		goto out;
+		return 0;
 	ret = edgetpu_mmu_attach_domain(group->etdev, group->etdomain);
 	if (ret)
-		goto out;
-	if (group->etdomain->pasid == IOMMU_PASID_INVALID)
-		mailbox = edgetpu_mailbox_vii_add(mgr, 0);
-	else
-		mailbox = edgetpu_mailbox_vii_add(mgr, group->etdomain->pasid);
-	if (IS_ERR(mailbox)) {
+		return ret;
+	ret = edgetpu_mailbox_init_vii(&group->vii, group, &group->mbox_attr);
+	if (ret) {
 		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
-		ret = PTR_ERR(mailbox);
-		goto out;
+		return ret;
 	}
-	ctx_id = mailbox->mailbox_id;
-	group->vii.mailbox = mailbox;
-	mailbox->internal.group = edgetpu_device_group_get(group);
-	edgetpu_mailbox_reinit_vii(group);
-	group->context_id = ctx_id;
-	if (edgetpu_device_group_is_finalized(group))
-		edgetpu_group_kci_open_device(group);
-out:
+	group->context_id = group->vii.mailbox->mailbox_id;
+	return 0;
+}
+
+int edgetpu_group_attach_and_open_mailbox(struct edgetpu_device_group *group)
+{
+	int ret;
+
+	mutex_lock(&group->lock);
+	ret = edgetpu_group_attach_mailbox_locked(group);
+	if (!ret && edgetpu_device_group_is_finalized(group))
+		ret = edgetpu_group_kci_open_device(group);
 	mutex_unlock(&group->lock);
 	return ret;
 }
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index cd47738..e0cd2a3 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -57,7 +57,7 @@ struct edgetpu_device_group {
 	uint workload_id;
 	struct edgetpu_dev *etdev;	/* the device opened by the leader */
 	/*
-	 * Whether edgetpu_group_detach_mailbox() has effects on this group.
+	 * Whether mailbox attaching and detaching have effects on this group.
 	 * This field is configured according to the priority field when
 	 * creating this group.
 	 */
@@ -333,16 +333,35 @@ bool edgetpu_set_group_join_lockout(struct edgetpu_dev *etdev, bool lockout);
 void edgetpu_fatal_error_notify(struct edgetpu_dev *etdev);
 
 /*
- * Detach and release the mailbox of VII from @group.
+ * Detach and release the mailbox resources of VII from @group.
  * Some group operations would be disabled when a group has no mailbox attached.
+ *
+ * Caller holds @group->lock.
+ */
+void edgetpu_group_detach_mailbox_locked(struct edgetpu_device_group *group);
+/*
+ * Before detaching the mailbox, send CLOSE_DEVICE KCI that claims the mailbox
+ * is going to be unused.
+ *
+ * The KCI command is sent even when @group is configured as mailbox
+ * non-detachable.
  */
-void edgetpu_group_detach_mailbox(struct edgetpu_device_group *group);
+void edgetpu_group_close_and_detach_mailbox(struct edgetpu_device_group *group);
 /*
- * Request and attach a mailbox of VII to @group.
+ * Request and attach the mailbox resources of VII to @group.
  *
  * Return 0 on success.
+ *
+ * Caller holds @group->lock.
+ */
+int edgetpu_group_attach_mailbox_locked(struct edgetpu_device_group *group);
+/*
+ * After (successfully) attaching the mailbox, send OPEN_DEVICE KCI.
+ *
+ * The KCI command is sent even when @group is configured as mailbox
+ * non-detachable (because the mailbox was successfully "attached").
  */
-int edgetpu_group_attach_mailbox(struct edgetpu_device_group *group);
+int edgetpu_group_attach_and_open_mailbox(struct edgetpu_device_group *group);
 
 /*
  * Checks whether @group has mailbox detached.
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 31111eb..0a348bd 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -36,6 +36,7 @@
 #include "edgetpu-mapping.h"
 #include "edgetpu-pm.h"
 #include "edgetpu-telemetry.h"
+#include "edgetpu-wakelock.h"
 #include "edgetpu.h"
 
 #define CREATE_TRACE_POINTS
@@ -112,11 +113,7 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 		return 0;
 	etdev = client->etdev;
 
-	mutex_lock(&client->wakelock.lock);
-
-	wakelock_count = client->wakelock.req_count;
-	/* Set wakelock state to "released" */
-	client->wakelock.req_count = 0;
+	wakelock_count = edgetpu_wakelock_lock(client->wakelock);
 
 	/* HACK: Can't disband a group if the device is off, turn it on */
 	if (client->group && !wakelock_count) {
@@ -124,7 +121,7 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 		edgetpu_pm_get(etdev->pm);
 	}
 
-	mutex_unlock(&client->wakelock.lock);
+	edgetpu_wakelock_unlock(client->wakelock);
 
 	edgetpu_client_remove(client);
 
@@ -533,13 +530,12 @@ static int edgetpu_ioctl_tpu_timestamp(struct edgetpu_client *client,
 	u64 timestamp;
 	int ret = 0;
 
-	mutex_lock(&client->wakelock.lock);
-	if (!client->wakelock.req_count) {
-		mutex_unlock(&client->wakelock.lock);
+	if (!edgetpu_wakelock_lock(client->wakelock)) {
+		edgetpu_wakelock_unlock(client->wakelock);
 		ret = -EAGAIN;
 	} else {
 		timestamp = edgetpu_chip_tpu_timestamp(client->etdev);
-		mutex_unlock(&client->wakelock.lock);
+		edgetpu_wakelock_unlock(client->wakelock);
 		if (copy_to_user(argp, &timestamp, sizeof(*argp)))
 			ret = -EFAULT;
 	}
@@ -553,79 +549,66 @@ static bool edgetpu_ioctl_check_permissions(struct file *file, uint cmd)
 
 static int edgetpu_ioctl_release_wakelock(struct edgetpu_client *client)
 {
-	if (!client->etdev->pm)
-		return -ENODEV;
-
-	mutex_lock(&client->wakelock.lock);
-
-	/* Cannot release wakelock if client has active CSR mappings */
-	if (client->wakelock.csr_map_count) {
-		etdev_warn(
-			client->etdev,
-			"%s: refusing wakelock release with %u CSR mappings\n",
-			__func__, client->wakelock.csr_map_count);
-		mutex_unlock(&client->wakelock.lock);
-		return -EAGAIN;
-	}
+	int count;
 
-	/* Cannot release wakelock if it wasn't acquired */
-	if (!client->wakelock.req_count) {
-		etdev_warn(client->etdev, "%s: invalid wakelock release\n",
-			   __func__);
-		mutex_unlock(&client->wakelock.lock);
-		return -EINVAL;
+	edgetpu_wakelock_lock(client->wakelock);
+	/* when NO_WAKELOCK: count should be 1 so here is a no-op */
+	count = edgetpu_wakelock_release(client->wakelock);
+	if (count < 0) {
+		edgetpu_wakelock_unlock(client->wakelock);
+		return count;
 	}
-
-	client->wakelock.req_count--;
-	if (!client->wakelock.req_count) {
+	if (!count) {
 		mutex_lock(&client->group_lock);
 		if (client->group)
-			edgetpu_group_detach_mailbox(client->group);
+			edgetpu_group_close_and_detach_mailbox(client->group);
 		mutex_unlock(&client->group_lock);
+		edgetpu_pm_put(client->etdev->pm);
 	}
-	edgetpu_pm_put(client->etdev->pm);
-	etdev_dbg(client->etdev,
-		  "%s: wakelock req count = %u CSR map count = %u\n", __func__,
-		  client->wakelock.req_count, client->wakelock.csr_map_count);
-	mutex_unlock(&client->wakelock.lock);
+	edgetpu_wakelock_unlock(client->wakelock);
+	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
+		  count);
 	return 0;
 }
 
 static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 {
+	int count;
 	int ret;
 
-	if (!client->etdev->pm)
-		return -ENODEV;
-
-	mutex_lock(&client->wakelock.lock);
-
-	ret = edgetpu_pm_get(client->etdev->pm);
-
-	if (ret) {
-		etdev_warn(client->etdev, "%s: pm_get failed (%d)", __func__,
-			   ret);
-		goto out_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_wakelock_unlock(client->wakelock);
+		return count;
 	}
-	if (!client->wakelock.req_count) {
+	if (!count) {
+		ret = edgetpu_pm_get(client->etdev->pm);
+		if (ret) {
+			etdev_warn(client->etdev, "%s: pm_get failed (%d)",
+				   __func__, ret);
+			goto error_release;
+		}
 		mutex_lock(&client->group_lock);
 		if (client->group)
-			ret = edgetpu_group_attach_mailbox(client->group);
+			ret = edgetpu_group_attach_and_open_mailbox(
+				client->group);
 		mutex_unlock(&client->group_lock);
 		if (ret) {
 			etdev_warn(client->etdev,
 				   "failed to attach mailbox: %d", ret);
 			edgetpu_pm_put(client->etdev->pm);
-			goto out_unlock;
+			goto error_release;
 		}
 	}
-
-	client->wakelock.req_count++;
-	etdev_dbg(client->etdev,
-		  "%s: wakelock req count = %u CSR map count = %u\n", __func__,
-		  client->wakelock.req_count, client->wakelock.csr_map_count);
-out_unlock:
-	mutex_unlock(&client->wakelock.lock);
+	edgetpu_wakelock_unlock(client->wakelock);
+	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
+		  count + 1);
+	return 0;
+error_release:
+	edgetpu_wakelock_release(client->wakelock);
+	edgetpu_wakelock_unlock(client->wakelock);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 37adebc..80691b0 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -99,6 +99,7 @@ struct edgetpu_reg_window {
 struct edgetpu_device_group;
 struct edgetpu_p2p_csr_map;
 struct edgetpu_remote_dram_map;
+struct edgetpu_wakelock;
 
 struct edgetpu_client {
 	pid_t pid;
@@ -126,11 +127,7 @@ struct edgetpu_client {
 	/* range of device CSRs mmap()'able */
 	struct edgetpu_reg_window reg_window;
 	/* Per-client request to keep device active */
-	struct {
-		struct mutex lock;
-		uint req_count;
-		uint csr_map_count;
-	} wakelock;
+	struct edgetpu_wakelock *wakelock;
 };
 
 struct edgetpu_mapping;
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 1a06578..c9061ec 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -712,23 +712,11 @@ void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_list_group *l;
 	struct edgetpu_device_group *group;
-	u32 mailbox_ids = 0;
 
 	mutex_lock(&etdev->groups_lock);
 	etdev_for_each_group(etdev, l, group) {
-		if (!edgetpu_group_mailbox_detached_locked(group)) {
+		if (!edgetpu_group_mailbox_detached_locked(group))
 			edgetpu_mailbox_reinit_vii(group);
-			if (edgetpu_device_group_is_finalized(group))
-				mailbox_ids |=
-					BIT(group->vii.mailbox->mailbox_id);
-		}
 	}
 	mutex_unlock(&etdev->groups_lock);
-	/*
-	 * If unfortunately groups are disbanded before we send this KCI, the
-	 * firmware side would be incorrectly informed that some mailboxes are
-	 * in use while actually not - but this shouldn't be harmful.
-	 */
-	if (mailbox_ids)
-		edgetpu_kci_open_device(etdev->kci, mailbox_ids);
 }
diff --git a/drivers/edgetpu/edgetpu-wakelock.c b/drivers/edgetpu/edgetpu-wakelock.c
new file mode 100644
index 0000000..ab9b919
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-wakelock.c
@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Wakelock for the runtime to explicitly claim it's going to use the EdgeTPU
+ * device.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+
+#include <linux/err.h>
+#include <linux/errno.h>
+#include <linux/mutex.h>
+#include <linux/slab.h>
+
+#include "edgetpu-config.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-wakelock.h"
+
+/*
+ * Returns the first event with a non-zero counter.
+ * Returns EDGETPU_WAKELOCK_EVENT_END if all event counters are zero.
+ *
+ * Caller holds @wakelock->lock.
+ */
+static enum edgetpu_wakelock_event
+wakelock_non_zero_event(struct edgetpu_wakelock *wakelock)
+{
+	int i;
+
+	for (i = 0; i < EDGETPU_WAKELOCK_EVENT_END; i++)
+		if (wakelock->event_count[i])
+			return i;
+	return EDGETPU_WAKELOCK_EVENT_END;
+}
+
+struct edgetpu_wakelock *edgetpu_wakelock_alloc(struct edgetpu_dev *etdev)
+{
+#ifndef EDGETPU_HAS_WAKELOCK
+	return EDGETPU_NO_WAKELOCK;
+#else /* !EDGETPU_HAS_WAKELOCK */
+	struct edgetpu_wakelock *wakelock =
+		kzalloc(sizeof(*wakelock), GFP_KERNEL);
+
+	if (!wakelock)
+		return NULL;
+	wakelock->etdev = etdev;
+	mutex_init(&wakelock->lock);
+	/* TODO(b/180528998): init as "released" */
+	/* Initialize client wakelock state to "acquired" */
+	wakelock->req_count = 1;
+	return wakelock;
+#endif /* EDGETPU_HAS_WAKELOCK */
+}
+
+void edgetpu_wakelock_free(struct edgetpu_wakelock *wakelock)
+{
+	if (IS_ERR_OR_NULL(wakelock))
+		return;
+	kfree(wakelock);
+}
+
+bool edgetpu_wakelock_inc_event(struct edgetpu_wakelock *wakelock,
+				enum edgetpu_wakelock_event evt)
+{
+	bool ret = true;
+
+	if (NO_WAKELOCK(wakelock))
+		return true;
+	mutex_lock(&wakelock->lock);
+	if (!wakelock->req_count) {
+		ret = false;
+		etdev_warn(
+			wakelock->etdev,
+			"invalid increase event %d when wakelock is released",
+			evt);
+	} else {
+		++wakelock->event_count[evt];
+		/* integer overflow.. */
+		if (unlikely(wakelock->event_count[evt] == 0)) {
+			--wakelock->event_count[evt];
+			ret = false;
+			etdev_warn_once(wakelock->etdev,
+					"int overflow on increasing event %d",
+					evt);
+		}
+	}
+	mutex_unlock(&wakelock->lock);
+	return ret;
+}
+
+bool edgetpu_wakelock_dec_event(struct edgetpu_wakelock *wakelock,
+				enum edgetpu_wakelock_event evt)
+{
+	bool ret = true;
+
+	if (NO_WAKELOCK(wakelock))
+		return true;
+	mutex_lock(&wakelock->lock);
+	if (!wakelock->event_count[evt]) {
+		ret = false;
+		etdev_warn(wakelock->etdev, "event %d unbalanced decreasing",
+			   evt);
+	} else {
+		--wakelock->event_count[evt];
+	}
+	mutex_unlock(&wakelock->lock);
+	return ret;
+}
+
+uint edgetpu_wakelock_lock(struct edgetpu_wakelock *wakelock)
+{
+	if (NO_WAKELOCK(wakelock))
+		return 1;
+	mutex_lock(&wakelock->lock);
+	return wakelock->req_count;
+}
+
+void edgetpu_wakelock_unlock(struct edgetpu_wakelock *wakelock)
+{
+	if (!NO_WAKELOCK(wakelock))
+		mutex_unlock(&wakelock->lock);
+}
+
+int edgetpu_wakelock_acquire(struct edgetpu_wakelock *wakelock)
+{
+	int ret;
+
+	if (NO_WAKELOCK(wakelock))
+		return 1;
+	ret = wakelock->req_count++;
+	/* integer overflow */
+	if (unlikely(ret < 0)) {
+		wakelock->req_count--;
+		return -EOVERFLOW;
+	}
+	return ret;
+}
+
+int edgetpu_wakelock_release(struct edgetpu_wakelock *wakelock)
+{
+	if (NO_WAKELOCK(wakelock))
+		return 1;
+	if (!wakelock->req_count) {
+		etdev_warn(wakelock->etdev, "invalid wakelock release");
+		return -EINVAL;
+	}
+	/* only need to check events when this is the last reference */
+	if (wakelock->req_count == 1) {
+		enum edgetpu_wakelock_event evt =
+			wakelock_non_zero_event(wakelock);
+
+		if (evt != EDGETPU_WAKELOCK_EVENT_END) {
+			etdev_warn(
+				wakelock->etdev,
+				"event %d is happening, refusing wakelock release",
+				evt);
+			return -EAGAIN;
+		}
+	}
+	return --wakelock->req_count;
+}
diff --git a/drivers/edgetpu/edgetpu-wakelock.h b/drivers/edgetpu/edgetpu-wakelock.h
new file mode 100644
index 0000000..9b6c812
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-wakelock.h
@@ -0,0 +1,139 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Wakelock for the runtime to explicitly claim it's going to use the EdgeTPU
+ * device.
+ *
+ * Copyright (C) 2021 Google, Inc.
+ */
+#ifndef __EDGETPU_WAKELOCK_H__
+#define __EDGETPU_WAKELOCK_H__
+
+#include <linux/err.h>
+#include <linux/mutex.h>
+
+#include "edgetpu-internal.h"
+
+/*
+ * See edgetpu_wakelock_alloc() for when this value is returned.
+ * Define as an errno so we can use IS_ERR* macros.
+ */
+#define EDGETPU_NO_WAKELOCK ERR_PTR(-EOPNOTSUPP)
+
+#define NO_WAKELOCK(wakelock) ((wakelock) == EDGETPU_NO_WAKELOCK)
+
+/*
+ * Events that could block the wakelock from being released.
+ * Use inc_event() and dec_event() to increase and decrease the counters when
+ * the event happens.
+ */
+enum edgetpu_wakelock_event {
+	EDGETPU_WAKELOCK_EVENT_FULL_CSR = 0,
+	EDGETPU_WAKELOCK_EVENT_MBOX_CSR = 1,
+	EDGETPU_WAKELOCK_EVENT_CMD_QUEUE = 2,
+	EDGETPU_WAKELOCK_EVENT_RESP_QUEUE = 3,
+
+	EDGETPU_WAKELOCK_EVENT_END = 4
+};
+
+struct edgetpu_wakelock {
+	struct edgetpu_dev *etdev; /* only for logging */
+	/* Protects every field below */
+	struct mutex lock;
+	/*
+	 * The request counter, increments on "acquire" and decrements on
+	 * "release".
+	 */
+	uint req_count;
+	/*
+	 * Events counter.
+	 * release() would fail if one of the slots is not zero.
+	 */
+	uint event_count[EDGETPU_WAKELOCK_EVENT_END];
+};
+
+/*
+ * Allocates and initializes a wakelock object.
+ *
+ * Returns the pointer on success, or NULL when out of memory.
+ *
+ * When the chipset doesn't support wakelock (judged by EDGETPU_HAS_WAKELOCK):
+ *   Nothing is allocated and this function returns EDGETPU_NO_WAKELOCK.
+ */
+struct edgetpu_wakelock *edgetpu_wakelock_alloc(struct edgetpu_dev *etdev);
+/* Frees the allocated wakelock. */
+void edgetpu_wakelock_free(struct edgetpu_wakelock *wakelock);
+
+/*
+ * Increases the event counter of @evt by one.
+ *
+ * Returns true if the counter is increased successfully.
+ * Returns false when one of the following errors happens:
+ *   - the wakelock is released
+ *   - integer overflow on the counter
+ *
+ * When the chipset doesn't support wakelock:
+ *   Does nothing and returns true.
+ */
+bool edgetpu_wakelock_inc_event(struct edgetpu_wakelock *wakelock,
+				enum edgetpu_wakelock_event evt);
+/*
+ * Decreases the event counter of @evt by one.
+ *
+ * Returns true if the counter is decreased successfully.
+ * Returns false when one of the following errors happens:
+ *   - the counter is zero
+ *
+ * When the chipset doesn't support wakelock:
+ *   Does nothing and returns true.
+ */
+bool edgetpu_wakelock_dec_event(struct edgetpu_wakelock *wakelock,
+				enum edgetpu_wakelock_event evt);
+
+/*
+ * Holds the internal lock of @wakelock. Fields in @wakelock are protected when
+ * this lock is holding.
+ *
+ * Returns the non-negative request counter of @wakelock.
+ *
+ * Example:
+ *   if (edgetpu_wakelock_lock(wakelock)) {
+ *      <..works that need the state of wakelock unchanged..>
+ *   }
+ *   edgetpu_wakelock_unlock(wakelock);
+ *
+ * When the chipset doesn't support wakelock:
+ *   Does nothing and returns 1.
+ */
+uint edgetpu_wakelock_lock(struct edgetpu_wakelock *wakelock);
+void edgetpu_wakelock_unlock(struct edgetpu_wakelock *wakelock);
+
+/*
+ * Acquires the wakelock, increases @wakelock->req_count by one.
+ *
+ * This function should be surrounded by edgetpu_wakelock_lock() and
+ * edgetpu_wakelock_unlock().
+ *
+ * Returns the value of request counter *before* being increased.
+ * Returns -EOVERFLOW if the request counter would overflow after increment.
+ *
+ * When the chipset doesn't support wakelock:
+ *   Does nothing and returns 1.
+ */
+int edgetpu_wakelock_acquire(struct edgetpu_wakelock *wakelock);
+/*
+ * Requests to release the wakelock, decreases @wakelock->req_count by one on
+ * success.
+ *
+ * This function should be surrounded by edgetpu_wakelock_lock() and
+ * edgetpu_wakelock_unlock().
+ *
+ * Returns the value of request counter *after* being decreased.
+ * Returns -EINVAL if the request counter is already zero.
+ * Returns -EAGAIN when there are events blocking wakelock from being released.
+ *
+ * When the chipset doesn't support wakelock:
+ *   Does nothing and returns 1.
+ */
+int edgetpu_wakelock_release(struct edgetpu_wakelock *wakelock);
+
+#endif /* __EDGETPU_WAKELOCK_H__ */