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

* whitechapel:
  edgetpu: send OPEN KCI only if wakelock acquired
  edgetpu: remove etdev state check for group join
  edgetpu: remove implicit wakelock on device open
  edgetpu: power on before group finalization
  edgetpu: simplify edgetpu_mailbox_init_vii
  edgetpu: init VII only when necessary
  edgetpu: move VII init from alloc to finalize
  edgetpu: add edgetpu_mailbox_validate_attr
  edgetpu: add snippets for mailbox at(de)taching
  edgetpu: code style fix on suspend/resume functions

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

7 files changed, 236 insertions, 178 deletions

diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 59bc79c..7510839 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -32,6 +32,7 @@
 #include "edgetpu-mmu.h"
 #include "edgetpu-sw-watchdog.h"
 #include "edgetpu-usr.h"
+#include "edgetpu-wakelock.h"
 #include "edgetpu.h"
 #include "mm-backport.h"
 
@@ -299,6 +300,45 @@ static void group_release_members(struct edgetpu_device_group *group)
 	group->members = NULL;
 }
 
+/*
+ * Does attach domain, init VII, and set @group->context_id without checking
+ * @group->mailbox_detachable and whether the mailbox is attached.
+ *
+ * Caller holds @group->lock.
+ */
+static int do_attach_mailbox_locked(struct edgetpu_device_group *group)
+{
+	int ret;
+
+	ret = edgetpu_mmu_attach_domain(group->etdev, group->etdomain);
+	if (ret)
+		return ret;
+	ret = edgetpu_mailbox_init_vii(&group->vii, group);
+	if (ret) {
+		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
+		return ret;
+	}
+	group->context_id = group->vii.mailbox->mailbox_id;
+	return 0;
+}
+
+/*
+ * Does detach domain, remove VII, and invalidate @group->context_id without
+ * checking @group->mailbox_detachable and whether the mailbox is detached.
+ *
+ * Caller holds @group->lock.
+ */
+static void do_detach_mailbox_locked(struct edgetpu_device_group *group)
+{
+	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;
+}
+
 int edgetpu_group_set_eventfd(struct edgetpu_device_group *group, uint event_id,
 			      int eventfd)
 {
@@ -380,9 +420,9 @@ static void edgetpu_device_group_release(struct edgetpu_device_group *group)
 #ifdef EDGETPU_HAS_P2P_MAILBOX
 		edgetpu_p2p_mailbox_release(group);
 #endif
+		edgetpu_mailbox_remove_vii(&group->vii);
 		group_release_members(group);
 	}
-	edgetpu_mailbox_remove_vii(&group->vii);
 	if (group->etdomain) {
 		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
 		edgetpu_mmu_free_domain(group->etdev, group->etdomain);
@@ -564,58 +604,6 @@ void edgetpu_device_group_leave_locked(struct edgetpu_client *client)
 	mutex_unlock(&client->etdev->groups_lock);
 }
 
-/* caller should hold client's etdev state lock. */
-static int edgetpu_device_group_add_locked(struct edgetpu_device_group *group,
-					   struct edgetpu_client *client)
-{
-	struct edgetpu_list_client *c;
-	int ret = 0;
-
-	mutex_lock(&client->group_lock);
-	if (client->group) {
-		mutex_unlock(&client->group_lock);
-		return -EINVAL;
-	}
-
-	mutex_lock(&group->lock);
-	if (!edgetpu_device_group_is_waiting(group)) {
-		ret = -EINVAL;
-		goto out;
-	}
-
-	for_each_list_client(c, group) {
-		if (!edgetpu_clients_groupable(c->client, client)) {
-			ret = -EINVAL;
-			goto out;
-		}
-	}
-
-	c = kzalloc(sizeof(*c), GFP_KERNEL);
-	if (!c) {
-		ret = -ENOMEM;
-		goto out;
-	}
-
-	ret = edgetpu_dev_add_group(client->etdev, group);
-	if (ret) {
-		kfree(c);
-		goto out;
-	}
-
-	c->client = edgetpu_client_get(client);
-	list_add_tail(&c->list, &group->clients);
-	client->idx = group->n_clients;
-	group->n_clients++;
-	client->group = edgetpu_device_group_get(group);
-	etdev_dbg(client->etdev, "%s: added group %u", __func__,
-		  group->workload_id);
-
-out:
-	mutex_unlock(&group->lock);
-	mutex_unlock(&client->group_lock);
-	return ret;
-}
-
 void edgetpu_device_group_leave(struct edgetpu_client *client)
 {
 	mutex_lock(&client->etdev->state_lock);
@@ -633,29 +621,24 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 			   const struct edgetpu_mailbox_attr *attr)
 {
 	static uint cur_workload_id;
-	int ret;
+	int ret = -EINVAL;
 	struct edgetpu_device_group *group;
 	struct edgetpu_iommu_domain *etdomain;
 
-	mutex_lock(&client->etdev->state_lock);
-	if (client->etdev->state != ETDEV_STATE_GOOD) {
-		ret = edgetpu_get_state_errno_locked(client->etdev);
-		goto state_unlock;
-	}
+	if (!edgetpu_mailbox_validate_attr(attr))
+		goto error;
 	/*
 	 * The client already belongs to a group.
 	 * It's safe not to take client->group_lock as
-	 * edgetpu_device_group_add_locked() will fail if there is race.
+	 * edgetpu_device_group_add() will fail if there is race.
 	 */
-	if (client->group) {
-		ret = -EINVAL;
-		goto state_unlock;
-	}
+	if (client->group)
+		goto error;
 
 	group = kzalloc(sizeof(*group), GFP_KERNEL);
 	if (!group) {
 		ret = -ENOMEM;
-		goto state_unlock;
+		goto error;
 	}
 
 	refcount_set(&group->ref_count, 1);
@@ -664,12 +647,16 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 	group->n_clients = 0;
 	group->status = EDGETPU_DEVICE_GROUP_WAITING;
 	group->etdev = client->etdev;
+	group->vii.etdev = client->etdev;
 	mutex_init(&group->lock);
 	rwlock_init(&group->events.lock);
 	edgetpu_mapping_init(&group->host_mappings);
 	edgetpu_mapping_init(&group->dmabuf_mappings);
+	group->mbox_attr = *attr;
+	if (attr->priority & EDGETPU_PRIORITY_DETACHABLE)
+		group->mailbox_detachable = true;
 	/* adds @client as the first entry */
-	ret = edgetpu_device_group_add_locked(group, client);
+	ret = edgetpu_device_group_add(group, client);
 	if (ret) {
 		etdev_dbg(group->etdev, "%s: group %u add failed ret=%d",
 			  __func__, group->workload_id, ret);
@@ -681,24 +668,12 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 		ret = -ENOMEM;
 		goto error_leave_group;
 	}
-	ret = edgetpu_mmu_attach_domain(group->etdev, etdomain);
-	if (ret) {
-		edgetpu_mmu_free_domain(group->etdev, etdomain);
-		goto error_leave_group;
-	}
 	group->etdomain = etdomain;
-	ret = edgetpu_mailbox_init_vii(&group->vii, group, attr);
-	if (ret) {
-		etdev_dbg(group->etdev, "%s: group %u init vii failed ret=%d",
-			  __func__, group->workload_id, ret);
-		/* this also performs domain detach / free */
-		goto error_leave_group;
-	}
-
-	group->context_id = group->vii.mailbox->mailbox_id;
-	if (attr->priority & EDGETPU_PRIORITY_DETACHABLE)
-		group->mailbox_detachable = true;
-	group->mbox_attr = *attr;
+	if (etdomain->token != EDGETPU_DOMAIN_TOKEN_END)
+		group->context_id =
+			EDGETPU_CONTEXT_DOMAIN_TOKEN | etdomain->token;
+	else
+		group->context_id = EDGETPU_CONTEXT_INVALID;
 
 	mutex_unlock(&client->etdev->state_lock);
 	return group;
@@ -707,24 +682,58 @@ error_leave_group:
 	edgetpu_device_group_leave_locked(client);
 error_put_group:
 	edgetpu_device_group_put(group);
-state_unlock:
-	mutex_unlock(&client->etdev->state_lock);
+error:
 	return ERR_PTR(ret);
 }
 
 int edgetpu_device_group_add(struct edgetpu_device_group *group,
 			     struct edgetpu_client *client)
 {
-	int ret;
+	struct edgetpu_list_client *c;
+	int ret = 0;
 
-	mutex_lock(&client->etdev->state_lock);
-	if (client->etdev->state != ETDEV_STATE_GOOD) {
-		ret = edgetpu_get_state_errno_locked(client->etdev);
+	mutex_lock(&client->group_lock);
+	if (client->group) {
+		mutex_unlock(&client->group_lock);
+		return -EINVAL;
+	}
+
+	mutex_lock(&group->lock);
+	if (!edgetpu_device_group_is_waiting(group)) {
+		ret = -EINVAL;
+		goto out;
+	}
+
+	for_each_list_client(c, group) {
+		if (!edgetpu_clients_groupable(c->client, client)) {
+			ret = -EINVAL;
+			goto out;
+		}
+	}
+
+	c = kzalloc(sizeof(*c), GFP_KERNEL);
+	if (!c) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	ret = edgetpu_dev_add_group(client->etdev, group);
+	if (ret) {
+		kfree(c);
 		goto out;
 	}
-	ret = edgetpu_device_group_add_locked(group, client);
+
+	c->client = edgetpu_client_get(client);
+	list_add_tail(&c->list, &group->clients);
+	client->idx = group->n_clients;
+	group->n_clients++;
+	client->group = edgetpu_device_group_get(group);
+	etdev_dbg(client->etdev, "%s: added group %u", __func__,
+		  group->workload_id);
+
 out:
-	mutex_unlock(&client->etdev->state_lock);
+	mutex_unlock(&group->lock);
+	mutex_unlock(&client->group_lock);
 	return ret;
 }
 
@@ -743,14 +752,15 @@ int edgetpu_device_group_finalize(struct edgetpu_device_group *group)
 {
 	int ret = 0, i;
 	struct edgetpu_dev *etdev;
+	bool mailbox_attached = false;
+	struct edgetpu_client *leader;
 
 	mutex_lock(&group->lock);
 	/* do nothing if the group is finalized */
 	if (edgetpu_device_group_is_finalized(group))
 		goto err_unlock;
 
-	if (!edgetpu_device_group_is_waiting(group) ||
-	    edgetpu_group_mailbox_detached_locked(group)) {
+	if (!edgetpu_device_group_is_waiting(group)) {
 		ret = -EINVAL;
 		goto err_unlock;
 	}
@@ -768,10 +778,29 @@ int edgetpu_device_group_finalize(struct edgetpu_device_group *group)
 		goto err_release_members;
 	}
 
+	leader = group->members[0];
+	/*
+	 * Initialize VII mailbox if
+	 * 1. mailbox is non-detachable: VII is assigned and has the same life
+	 *    cycle as a finalized @group, or
+	 * 2. has non-zero wakelock reference counter: VII should be ready to
+	 *    use after group is finalized.
+	 */
+	if (!group->mailbox_detachable ||
+	    edgetpu_wakelock_count_locked(leader->wakelock)) {
+		mailbox_attached = true;
+		ret = do_attach_mailbox_locked(group);
+		if (ret) {
+			etdev_err(group->etdev,
+				  "attach mailbox failed on finalization: %d",
+				  ret);
+			goto err_release_members;
+		}
+	}
 #ifdef EDGETPU_HAS_P2P_MAILBOX
 	ret = edgetpu_p2p_mailbox_setup(group);
 	if (ret)
-		goto err_release_members;
+		goto err_detach_mailbox;
 #endif
 
 	ret = edgetpu_group_setup_remote_dram(group);
@@ -779,9 +808,12 @@ int edgetpu_device_group_finalize(struct edgetpu_device_group *group)
 		goto err_release_p2p;
 
 	edgetpu_usr_init_group(group);
-	ret = edgetpu_device_group_kci_finalized(group);
-	if (ret)
-		goto err_remove_remote_dram;
+
+	if (edgetpu_wakelock_count_locked(leader->wakelock)) {
+		ret = edgetpu_device_group_kci_finalized(group);
+		if (ret)
+			goto err_remove_remote_dram;
+	}
 
 	group->status = EDGETPU_DEVICE_GROUP_FINALIZED;
 
@@ -798,7 +830,10 @@ err_remove_remote_dram:
 err_release_p2p:
 #ifdef EDGETPU_HAS_P2P_MAILBOX
 	edgetpu_p2p_mailbox_release(group);
+err_detach_mailbox:
 #endif
+	if (mailbox_attached)
+		do_detach_mailbox_locked(group);
 err_release_members:
 	group_release_members(group);
 err_unlock:
@@ -1532,52 +1567,40 @@ void edgetpu_group_detach_mailbox_locked(struct edgetpu_device_group *group)
 		return;
 	if (edgetpu_group_mailbox_detached_locked(group))
 		return;
-	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;
+	do_detach_mailbox_locked(group);
 }
 
 void edgetpu_group_close_and_detach_mailbox(struct edgetpu_device_group *group)
 {
 	mutex_lock(&group->lock);
-	if (edgetpu_device_group_is_finalized(group))
+	/* only a finalized group may have mailbox attached */
+	if (edgetpu_device_group_is_finalized(group)) {
 		edgetpu_group_kci_close_device(group);
-	edgetpu_group_detach_mailbox_locked(group);
+		edgetpu_group_detach_mailbox_locked(group);
+	}
 	mutex_unlock(&group->lock);
 }
 
 int edgetpu_group_attach_mailbox_locked(struct edgetpu_device_group *group)
 {
-	int ret;
-
 	if (!group->mailbox_detachable)
 		return 0;
 	if (!edgetpu_group_mailbox_detached_locked(group))
 		return 0;
-	ret = edgetpu_mmu_attach_domain(group->etdev, group->etdomain);
-	if (ret)
-		return ret;
-	ret = edgetpu_mailbox_init_vii(&group->vii, group, &group->mbox_attr);
-	if (ret) {
-		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
-		return ret;
-	}
-	group->context_id = group->vii.mailbox->mailbox_id;
-	return 0;
+	return do_attach_mailbox_locked(group);
 }
 
 int edgetpu_group_attach_and_open_mailbox(struct edgetpu_device_group *group)
 {
-	int ret;
+	int ret = 0;
 
 	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);
+	/* only attaching mailbox for finalized groups */
+	if (edgetpu_device_group_is_finalized(group)) {
+		ret = edgetpu_group_attach_mailbox_locked(group);
+		if (!ret)
+			ret = edgetpu_group_kci_open_device(group);
+	}
 	mutex_unlock(&group->lock);
 	return ret;
 }
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 0a348bd..9efee4b 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -69,21 +69,10 @@ static struct dentry *edgetpu_debugfs_dir;
 int edgetpu_open(struct edgetpu_dev *etdev, struct file *file)
 {
 	struct edgetpu_client *client;
-	int res;
 
 	/* Set client pointer to NULL if error creating client. */
 	file->private_data = NULL;
 	mutex_lock(&etdev->open.lock);
-	if (etdev->pm) {
-		res = edgetpu_pm_get(etdev->pm);
-		if (res) {
-			dev_err(etdev->dev,
-				"Failed to request device power up (%d)", res);
-			mutex_unlock(&etdev->open.lock);
-			return -ENODEV;
-		}
-	}
-
 	client = edgetpu_client_add(etdev);
 	if (IS_ERR(client)) {
 		mutex_unlock(&etdev->open.lock);
@@ -115,7 +104,10 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 
 	wakelock_count = edgetpu_wakelock_lock(client->wakelock);
 
-	/* HACK: Can't disband a group if the device is off, turn it on */
+	/*
+	 * TODO(b/180528495): remove pm_get when disbanding can be performed
+	 * with device off.
+	 */
 	if (client->group && !wakelock_count) {
 		wakelock_count = 1;
 		edgetpu_pm_get(etdev->pm);
@@ -205,11 +197,35 @@ static int edgetpu_ioctl_unset_perdie_eventfd(struct edgetpu_dev *etdev,
 static int edgetpu_ioctl_finalize_group(struct edgetpu_client *client)
 {
 	struct edgetpu_device_group *group;
-	int ret;
+	int ret = -EINVAL, wakelock_count;
 
-	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	/*
+	 * Hold the wakelock since we need to decide whether VII should be
+	 * initialized during finalization.
+	 */
+	wakelock_count = edgetpu_wakelock_lock(client->wakelock);
+	LOCK(client);
+	group = client->group;
+	if (!group || !edgetpu_device_group_is_leader(group, client))
+		goto out_unlock;
+	/*
+	 * TODO(b/180528495): remove pm_get when finalization can be performed
+	 * with device off.
+	 */
+	if (!wakelock_count) {
+		ret = edgetpu_pm_get(client->etdev->pm);
+		if (ret) {
+			etdev_err(client->etdev, "%s: pm_get failed (%d)",
+				  __func__, ret);
+			goto out_unlock;
+		}
+	}
 	ret = edgetpu_device_group_finalize(group);
+	if (!wakelock_count)
+		edgetpu_pm_put(client->etdev->pm);
+out_unlock:
 	UNLOCK(client);
+	edgetpu_wakelock_unlock(client->wakelock);
 	return ret;
 }
 
@@ -928,7 +944,7 @@ int edgetpu_fs_add(struct edgetpu_dev *etdev)
 	}
 
 	etdev->etcdev = device_create(edgetpu_class, etdev->dev, etdev->devno,
-				      etdev, etdev->dev_name);
+				      etdev, "%s", etdev->dev_name);
 	if (IS_ERR(etdev->etcdev)) {
 		ret = PTR_ERR(etdev->etcdev);
 		dev_err(etdev->dev, "%s: failed to create char device: %d\n",
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index c9061ec..3bdc7c5 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -392,21 +392,28 @@ static int convert_runtime_queue_size_to_fw(u32 queue_size, u32 element_size)
 	return ret;
 }
 
-/*
- * Sets mailbox and allocates queues to @vii.
- *
- * @group is the device group that @vii will be associated with.
- *
- * Returns 0 on success.
- * Returns -EINVAL if any fields in @attr is invalid.
- */
+bool edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr)
+{
+	int size;
+
+	size = convert_runtime_queue_size_to_fw(attr->cmd_queue_size,
+						attr->sizeof_cmd);
+	if (size < 0)
+		return false;
+	size = convert_runtime_queue_size_to_fw(attr->resp_queue_size,
+						attr->sizeof_resp);
+	if (size < 0)
+		return false;
+	return true;
+}
+
 int edgetpu_mailbox_init_vii(struct edgetpu_vii *vii,
-			     struct edgetpu_device_group *group,
-			     const struct edgetpu_mailbox_attr *attr)
+			     struct edgetpu_device_group *group)
 {
 	int cmd_queue_size, resp_queue_size;
 	struct edgetpu_mailbox_manager *mgr = group->etdev->mailbox_manager;
 	struct edgetpu_mailbox *mailbox;
+	const struct edgetpu_mailbox_attr *attr = &group->mbox_attr;
 	int ret;
 
 	if (!group->etdomain || group->etdomain->pasid == IOMMU_PASID_INVALID)
@@ -418,16 +425,8 @@ int edgetpu_mailbox_init_vii(struct edgetpu_vii *vii,
 
 	cmd_queue_size = convert_runtime_queue_size_to_fw(attr->cmd_queue_size,
 							  attr->sizeof_cmd);
-	if (cmd_queue_size < 0) {
-		edgetpu_mailbox_remove(mgr, mailbox);
-		return cmd_queue_size;
-	}
 	resp_queue_size = convert_runtime_queue_size_to_fw(
 		attr->resp_queue_size, attr->sizeof_resp);
-	if (resp_queue_size < 0) {
-		edgetpu_mailbox_remove(mgr, mailbox);
-		return resp_queue_size;
-	}
 
 	edgetpu_mailbox_set_priority(mailbox, attr->priority);
 	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox,
@@ -674,12 +673,8 @@ void edgetpu_mailbox_reinit_vii(struct edgetpu_device_group *group)
 {
 	int cmd_queue_size, resp_queue_size;
 	struct edgetpu_mailbox *mailbox = group->vii.mailbox;
-	struct edgetpu_mailbox_attr *attr = &group->mbox_attr;
+	const struct edgetpu_mailbox_attr *attr = &group->mbox_attr;
 
-	/*
-	 * Sizes here should never be invalid since they are checked in
-	 * edgetpu_mailbox_init_vii().
-	 */
 	cmd_queue_size = convert_runtime_queue_size_to_fw(attr->cmd_queue_size,
 							  attr->sizeof_cmd);
 	resp_queue_size = convert_runtime_queue_size_to_fw(
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index 238915f..026bdbf 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -214,9 +214,27 @@ void edgetpu_mailbox_inc_resp_queue_head(struct edgetpu_mailbox *mailbox,
  */
 struct edgetpu_mailbox *
 edgetpu_mailbox_vii_add(struct edgetpu_mailbox_manager *mgr, uint id);
+/*
+ * Validates the mailbox attributes.
+ * Returns true if valid, false otherwise.
+ *
+ * See the error cases of EDGETPU_CREATE_GROUP in edgetpu.h for when will @attr
+ * be considered as invalid.
+ */
+bool edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr);
+/*
+ * Sets mailbox and allocates queues to @vii.
+ *
+ * @group is the device group that @vii will be associated with,
+ * @group->mbox_attr is used to set the VII mailbox attributes.
+ *
+ * @group->mbox_attr must be checked by edgetpu_mailbox_validate_attr() before
+ * calling this function.
+ *
+ * Returns 0 on success.
+ */
 int edgetpu_mailbox_init_vii(struct edgetpu_vii *vii,
-			     struct edgetpu_device_group *group,
-			     const struct edgetpu_mailbox_attr *attr);
+			     struct edgetpu_device_group *group);
 void edgetpu_mailbox_remove_vii(struct edgetpu_vii *vii);
 
 
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index c29ce9e..8ee47fe 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -259,20 +259,16 @@ void edgetpu_pchannel_power_up(struct edgetpu_dev *etdev)
 	pchannel_state_change_request(etdev, STATE_RUN);
 }
 
-
 #if IS_ENABLED(CONFIG_PM_SLEEP)
 
 int edgetpu_pm_suspend(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_pm *etpm = etdev->pm;
 
-	if (!etpm)
-		return 0;
-
-	if (etpm->p->power_up_count) {
+	if (etpm && etpm->p->power_up_count) {
 		etdev_warn_ratelimited(
-			etdev, "%s: cannot suspend with power up count = %d\n",
-			__func__, etpm->p->power_up_count);
+			etdev, "cannot suspend with power up count = %d\n",
+			etpm->p->power_up_count);
 		return -EAGAIN;
 	}
 
@@ -283,13 +279,10 @@ int edgetpu_pm_resume(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_pm *etpm = etdev->pm;
 
-	if (!etpm)
-		return 0;
-
-	if (etpm->p->power_up_count)
+	if (etpm && etpm->p->power_up_count)
 		etdev_warn_ratelimited(etdev,
-				       "%s: resumed with power up count = %d\n",
-				       __func__, etpm->p->power_up_count);
+				       "resumed with power up count = %d\n",
+				       etpm->p->power_up_count);
 
 	return 0;
 }
diff --git a/drivers/edgetpu/edgetpu-wakelock.c b/drivers/edgetpu/edgetpu-wakelock.c
index ab9b919..61af4fa 100644
--- a/drivers/edgetpu/edgetpu-wakelock.c
+++ b/drivers/edgetpu/edgetpu-wakelock.c
@@ -44,9 +44,8 @@ struct edgetpu_wakelock *edgetpu_wakelock_alloc(struct edgetpu_dev *etdev)
 		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;
+	/* Initialize client wakelock state to "released" */
+	wakelock->req_count = 0;
 	return wakelock;
 #endif /* EDGETPU_HAS_WAKELOCK */
 }
diff --git a/drivers/edgetpu/edgetpu-wakelock.h b/drivers/edgetpu/edgetpu-wakelock.h
index 9b6c812..713f58b 100644
--- a/drivers/edgetpu/edgetpu-wakelock.h
+++ b/drivers/edgetpu/edgetpu-wakelock.h
@@ -108,6 +108,20 @@ uint edgetpu_wakelock_lock(struct edgetpu_wakelock *wakelock);
 void edgetpu_wakelock_unlock(struct edgetpu_wakelock *wakelock);
 
 /*
+ * Returns the request counter of @wakelock.
+ *
+ * Caller calls edgetpu_wakelock_lock() before calling this function.
+ *
+ * When the chipset doesn't support wakelock:
+ *   Returns 1.
+ */
+static inline uint
+edgetpu_wakelock_count_locked(struct edgetpu_wakelock *wakelock)
+{
+	return NO_WAKELOCK(wakelock) ? 1 : wakelock->req_count;
+}
+
+/*
  * Acquires the wakelock, increases @wakelock->req_count by one.
  *
  * This function should be surrounded by edgetpu_wakelock_lock() and