[Copybara Auto Merge] Merge branch zuma into android14-gs-pixel-5.15

gcip: fixup: Use devm_* for gcip-pm
Bug: 265870718
edgetpu: Use const gcip_fw_tracing_args
Bug: 262916889
gcip: Use devm_* for gcip-pm
Bug: 265870718 (repeat)
gcip: Remove the start log of async power down
Bug: 265870718 (repeat)
gcip: Add a comment about suspend/resume in gcip-pm
Bug: 265870718 (repeat)
gcip: Add firmware dynamic tracing support
Bug: 262916889 (repeat)
gcip: Add PM support
Bug: 265870718 (repeat)
gcip: call before_enqueue_wait_list
Bug: 267713283
gcip: Add firmware dynamic tracing header
Bug: 262916889 (repeat)
gcip: Add PM header
Bug: 265870718 (repeat)
gcip: introduce before_enqueue_wait_list
Bug: 267713283 (repeat)
edgetpu: Adopt GCIP fw tracing
Bug: 262916889 (repeat)
edgetpu: Adopt GCIP PM
Bug: 265870718 (repeat)

Signed-off-by: Zuma copybara merger <zuma-automerger@google.com>
GitOrigin-RevId: 961e22cac86af77c8609fbe7c4d08d89e4aea8b3
Change-Id: Ia27beb7af534208da90b1bcecdf7f3478bf3f721

23 files changed, 763 insertions, 392 deletions

diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 8bf8ab5..fa0f849 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -22,6 +22,8 @@
 #include <linux/types.h>
 #include <linux/uidgid.h>
 
+#include <gcip/gcip-firmware.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-debug-dump.h"
 #include "edgetpu-device-group.h"
@@ -29,6 +31,7 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 #include "edgetpu-mmu.h"
+#include "edgetpu-pm.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-sw-watchdog.h"
 #include "edgetpu-telemetry.h"
@@ -392,6 +395,24 @@ int edgetpu_get_state_errno_locked(struct edgetpu_dev *etdev)
 	return 0;
 }
 
+static struct gcip_fw_tracing *edgetpu_firmware_tracing_create(struct edgetpu_dev *etdev)
+{
+	const struct gcip_fw_tracing_args fw_tracing_args = {
+		.dev = etdev->dev,
+		.pm = etdev->pm,
+		.dentry = edgetpu_fs_debugfs_dir(),
+		.data = etdev,
+		.set_level = edgetpu_kci_firmware_tracing_level,
+	};
+
+	return gcip_firmware_tracing_create(&fw_tracing_args);
+}
+
+static void edgetpu_firmware_tracing_destroy(struct gcip_fw_tracing *fw_tracing)
+{
+	gcip_firmware_tracing_destroy(fw_tracing);
+}
+
 int edgetpu_device_add(struct edgetpu_dev *etdev,
 		       const struct edgetpu_mapped_resource *regs,
 		       const struct edgetpu_iface_params *iface_params,
@@ -487,6 +508,12 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 	if (ret)
 		etdev_warn(etdev, "debug dump init fail: %d", ret);
 
+	etdev->fw_tracing = edgetpu_firmware_tracing_create(etdev);
+	if (IS_ERR(etdev->fw_tracing)) {
+		etdev_warn(etdev, "firmware tracing create fail: %ld", PTR_ERR(etdev->fw_tracing));
+		etdev->fw_tracing = NULL;
+	}
+
 	edgetpu_chip_init(etdev);
 	/* No limit on DMA segment size */
 	dma_set_max_seg_size(etdev->dev, UINT_MAX);
@@ -506,6 +533,7 @@ remove_dev:
 void edgetpu_device_remove(struct edgetpu_dev *etdev)
 {
 	edgetpu_chip_exit(etdev);
+	edgetpu_firmware_tracing_destroy(etdev->fw_tracing);
 	edgetpu_debug_dump_exit(etdev);
 	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
 	edgetpu_usage_stats_exit(etdev);
diff --git a/drivers/edgetpu/edgetpu-debug-dump.c b/drivers/edgetpu/edgetpu-debug-dump.c
index f866488..a04dcbf 100644
--- a/drivers/edgetpu/edgetpu-debug-dump.c
+++ b/drivers/edgetpu/edgetpu-debug-dump.c
@@ -9,17 +9,18 @@
 #include <linux/debugfs.h>
 #include <linux/workqueue.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-debug-dump.h"
 #include "edgetpu-device-group.h"
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-kci.h"
-#include "edgetpu-pm.h"
 
 static int edgetpu_get_debug_dump_set(void *data, u64 val)
 {
 	struct edgetpu_dev *etdev = data;
-	int ret = edgetpu_pm_get(etdev->pm);
+	int ret = gcip_pm_get(etdev->pm);
 
 	if (ret)
 		return ret;
@@ -28,7 +29,7 @@ static int edgetpu_get_debug_dump_set(void *data, u64 val)
 		etdev_warn(etdev, "FW refused debug dump request: %d", ret);
 		ret = -EOPNOTSUPP;
 	}
-	edgetpu_pm_put(etdev->pm);
+	gcip_pm_put(etdev->pm);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 008b92a..6bfc199 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -15,6 +15,8 @@
 #include <linux/string.h>
 #include <linux/types.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu.h"
 #include "edgetpu-debug-dump.h"
 #include "edgetpu-device-group.h"
@@ -22,7 +24,6 @@
 #include "edgetpu-firmware-util.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-kci.h"
-#include "edgetpu-pm.h"
 #include "edgetpu-sw-watchdog.h"
 #include "edgetpu-telemetry.h"
 
@@ -177,6 +178,10 @@ static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 	if (ret)
 		etdev_warn(etdev, "telemetry KCI error: %d", ret);
 
+	ret = gcip_firmware_tracing_restore(etdev->fw_tracing);
+	if (ret)
+		etdev_warn(etdev, "firmware tracing restore error: %d", ret);
+
 	ret = edgetpu_thermal_restore(etdev);
 	if (ret)
 		etdev_warn(etdev, "thermal restore error: %d", ret);
@@ -187,9 +192,9 @@ static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 }
 
 /*
- * Do edgetpu_pm_get() but prevent it from running the loaded firmware.
+ * Do gcip_pm_get() but prevent it from running the loaded firmware.
  *
- * On success, caller must later call edgetpu_pm_put() to decrease the reference count.
+ * On success, caller must later call gcip_pm_put() to decrease the reference count.
  *
  * Caller holds firmware lock.
  */
@@ -201,7 +206,7 @@ static int edgetpu_firmware_pm_get(struct edgetpu_firmware *et_fw)
 	/* Prevent platform-specific code from trying to run the previous firmware */
 	et_fw->p->status = GCIP_FW_LOADING;
 	etdev_dbg(et_fw->etdev, "Requesting power up for firmware run\n");
-	ret = edgetpu_pm_get(et_fw->etdev->pm);
+	ret = gcip_pm_get(et_fw->etdev->pm);
 	if (ret)
 		et_fw->p->status = prev;
 	return ret;
@@ -388,7 +393,7 @@ int edgetpu_firmware_run(struct edgetpu_dev *etdev, const char *name,
 	ret = edgetpu_firmware_pm_get(et_fw);
 	if (!ret) {
 		ret = edgetpu_firmware_run_locked(et_fw, name, flags);
-		edgetpu_pm_put(etdev->pm);
+		gcip_pm_put(etdev->pm);
 	}
 
 	edgetpu_firmware_load_unlock(etdev);
@@ -597,7 +602,7 @@ static void edgetpu_firmware_wdt_timeout_action(void *data)
 
 	etdev->watchdog_timeout_count++;
 	/* Don't attempt f/w restart if device is off. */
-	if (!edgetpu_is_powered(etdev))
+	if (!gcip_pm_is_powered(etdev->pm))
 		return;
 
 	/*
@@ -617,7 +622,7 @@ static void edgetpu_firmware_wdt_timeout_action(void *data)
 	ret = edgetpu_firmware_pm_get(et_fw);
 	if (!ret) {
 		ret = edgetpu_firmware_restart_locked(etdev, true);
-		edgetpu_pm_put(etdev->pm);
+		gcip_pm_put(etdev->pm);
 	}
 	edgetpu_firmware_unlock(etdev);
 }
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 43b734f..4cc1e70 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -29,6 +29,8 @@
 #include <linux/uaccess.h>
 #include <linux/uidgid.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-device-group.h"
 #include "edgetpu-dmabuf.h"
@@ -36,7 +38,6 @@
 #include "edgetpu-internal.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mapping.h"
-#include "edgetpu-pm.h"
 #include "edgetpu-telemetry.h"
 #include "edgetpu-wakelock.h"
 #include "edgetpu.h"
@@ -124,7 +125,7 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 
 	/* count was zero if client previously released its wake lock */
 	if (wakelock_count)
-		edgetpu_pm_put(etdev->pm);
+		gcip_pm_put(etdev->pm);
 	return 0;
 }
 
@@ -214,7 +215,7 @@ static int edgetpu_ioctl_finalize_group(struct edgetpu_client *client)
 	if (!group)
 		goto out_unlock;
 	/* Finalization has to be performed with device on. */
-	ret = edgetpu_pm_get(client->etdev->pm);
+	ret = gcip_pm_get(client->etdev->pm);
 	if (ret) {
 		etdev_err(client->etdev, "%s: pm_get failed (%d)",
 			  __func__, ret);
@@ -227,7 +228,7 @@ static int edgetpu_ioctl_finalize_group(struct edgetpu_client *client)
 	edgetpu_wakelock_lock(client->wakelock);
 	ret = edgetpu_device_group_finalize(group);
 	edgetpu_wakelock_unlock(client->wakelock);
-	edgetpu_pm_put(client->etdev->pm);
+	gcip_pm_put(client->etdev->pm);
 out_unlock:
 	UNLOCK(client);
 	return ret;
@@ -483,7 +484,7 @@ static int edgetpu_ioctl_release_wakelock(struct edgetpu_client *client)
 	if (!count) {
 		if (client->group)
 			edgetpu_group_close_and_detach_mailbox(client->group);
-		edgetpu_pm_put(client->etdev->pm);
+		gcip_pm_put(client->etdev->pm);
 	}
 	edgetpu_wakelock_unlock(client->wakelock);
 	UNLOCK(client);
@@ -524,7 +525,7 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 		goto error_client_unlock;
 	}
 
-	ret = edgetpu_pm_get(client->etdev->pm);
+	ret = gcip_pm_get(client->etdev->pm);
 	if (ret) {
 		etdev_warn(client->etdev, "pm_get failed (%d)", ret);
 		goto error_client_unlock;
@@ -551,7 +552,7 @@ error_wakelock_unlock:
 
 	/* Balance the power up count due to pm_get above.*/
 	if (ret || count)
-		edgetpu_pm_put(client->etdev->pm);
+		gcip_pm_put(client->etdev->pm);
 
 error_client_unlock:
 	UNLOCK(client);
@@ -807,9 +808,9 @@ static int edgetpu_pm_debugfs_set_wakelock(void *data, u64 val)
 	int ret = 0;
 
 	if (val)
-		ret = edgetpu_pm_get(etdev->pm);
+		ret = gcip_pm_get(etdev->pm);
 	else
-		edgetpu_pm_put(etdev->pm);
+		gcip_pm_put(etdev->pm);
 	return ret;
 }
 DEFINE_DEBUGFS_ATTRIBUTE(fops_wakelock, NULL, edgetpu_pm_debugfs_set_wakelock,
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 7c7a433..798d2d5 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -30,8 +30,10 @@
 #include <linux/types.h>
 #include <linux/workqueue.h>
 
+#include <gcip/gcip-firmware.h>
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu.h"
-#include "edgetpu-pm.h"
 #include "edgetpu-thermal.h"
 #include "edgetpu-usage-stats.h"
 
@@ -196,10 +198,11 @@ struct edgetpu_dev {
 	struct edgetpu_mailbox_manager *mailbox_manager;
 	struct edgetpu_kci *etkci;
 	struct edgetpu_firmware *firmware; /* firmware management */
+	struct gcip_fw_tracing *fw_tracing; /* firmware tracing */
 	struct edgetpu_telemetry_ctx *telemetry;
 	struct edgetpu_thermal *thermal;
 	struct edgetpu_usage_stats *usage_stats; /* usage stats private data */
-	struct edgetpu_pm *pm;  /* Power management interface */
+	struct gcip_pm *pm; /* Power management interface */
 	/* Memory pool in instruction remap region */
 	struct edgetpu_mempool *iremap_pool;
 	struct edgetpu_sw_wdt *etdev_sw_wdt;	/* software watchdog */
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index ed9b9e2..77c56f9 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -15,6 +15,7 @@
 #include <linux/slab.h>
 #include <linux/string.h> /* memcpy */
 
+#include <gcip/gcip-pm.h>
 #include <gcip/gcip-telemetry.h>
 
 #include "edgetpu-firmware.h"
@@ -440,8 +441,9 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 	int ret = -EAGAIN;
 
 	/* Quick return if device is already powered down. */
-	if (!edgetpu_is_powered(etdev))
+	if (!gcip_pm_is_powered(etdev->pm))
 		return -EAGAIN;
+
 	/*
 	 * Lockout change in f/w load/unload status during usage update.
 	 * Skip usage update if the firmware is being updated now or is not
@@ -452,17 +454,12 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 
 	if (edgetpu_firmware_status_locked(etdev) != GCIP_FW_VALID)
 		goto fw_unlock;
-	/*
-	 * This function may run in a worker that is being canceled when the
-	 * device is powering down, and the power down code holds the PM lock.
-	 * Using trylock to prevent cancel_work_sync() waiting forever.
-	 */
-	if (!edgetpu_pm_trylock(etdev->pm))
+
+	if (gcip_pm_get_if_powered(etdev->pm, false))
 		goto fw_unlock;
 
-	if (edgetpu_is_powered(etdev))
-		ret = edgetpu_kci_update_usage_locked(etdev);
-	edgetpu_pm_unlock(etdev->pm);
+	ret = edgetpu_kci_update_usage_locked(etdev);
+	gcip_pm_put_async(etdev->pm);
 
 fw_unlock:
 	edgetpu_firmware_unlock(etdev);
@@ -636,18 +633,37 @@ int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block)
 	return gcip_kci_send_cmd(etdev->etkci->kci, &cmd);
 }
 
-int edgetpu_kci_thermal_control(struct edgetpu_dev *etdev, bool enable)
+int edgetpu_kci_firmware_tracing_level(void *data, unsigned long level, unsigned long *active_level)
 {
+	struct edgetpu_dev *etdev = data;
 	struct gcip_kci_command_element cmd = {
-		.code = GCIP_KCI_CODE_THERMAL_CONTROL,
+		.code = GCIP_KCI_CODE_FIRMWARE_TRACING_LEVEL,
 		.dma = {
-			.flags = (u32)enable,
+			.flags = (u32)level,
 		},
 	};
+	struct gcip_kci_response_element resp;
+	int ret;
 
 	if (!etdev->etkci)
 		return -ENODEV;
 
+	ret = gcip_kci_send_cmd_return_resp(etdev->etkci->kci, &cmd, &resp);
+	if (ret == GCIP_KCI_ERROR_OK)
+		*active_level = resp.retval;
+
+	return ret;
+}
+
+int edgetpu_kci_thermal_control(struct edgetpu_dev *etdev, bool enable)
+{
+	struct gcip_kci_command_element cmd = {
+		.code = GCIP_KCI_CODE_THERMAL_CONTROL,
+		.dma = {
+			.flags = (u32)enable,
+		},
+	};
+
 	return gcip_kci_send_cmd(etdev->etkci->kci, &cmd);
 }
 
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index 6b75306..ec9154e 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -125,7 +125,7 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev);
 /*
  * Works the same as edgetpu_kci_update_usage() except the caller of this
  * function must guarantee the device stays powered up, typically by calling
- * edgetpu_pm_get() or by calling this function from the power management
+ * gcip_pm_get() or by calling this function from the power management
  * functions themselves.
  *
  * Returns KCI response code on success or < 0 on error (typically -ETIMEDOUT).
@@ -196,6 +196,10 @@ 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(void *data, unsigned long level,
+				       unsigned long *active_level);
+
 /*
  * Request the firmware to enable or disable the thermal throttling.
  *
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 3bd4b2f..3dc02b4 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -14,6 +14,8 @@
 #include <linux/mmzone.h> /* MAX_ORDER_NR_PAGES */
 #include <linux/slab.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-device-group.h"
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-kci.h"
@@ -938,7 +940,12 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (edgetpu_pm_get_if_powered(group->etdev->pm, false)) {
+	if (gcip_pm_get_if_powered(group->etdev->pm, true)) {
+		mutex_lock(&group->lock);
+		ret = edgetpu_mailbox_external_alloc(group, req);
+		mutex_unlock(&group->lock);
+		goto out;
+	} else {
 		mutex_lock(&group->lock);
 		ret = edgetpu_mailbox_external_alloc(group, req);
 		if (ret) {
@@ -948,16 +955,11 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 		edgetpu_mailbox_init_external_mailbox(group->ext_mailbox);
 		ret = edgetpu_mailbox_activate_external_mailbox(group);
 		mutex_unlock(&group->lock);
-		edgetpu_pm_put(group->etdev->pm);
-		goto out;
-	} else {
-		mutex_lock(&group->lock);
-		ret = edgetpu_mailbox_external_alloc(group, req);
-		mutex_unlock(&group->lock);
+		gcip_pm_put(group->etdev->pm);
 		goto out;
 	}
 err:
-	edgetpu_pm_put(group->etdev->pm);
+	gcip_pm_put(group->etdev->pm);
 out:
 	edgetpu_device_group_put(group);
 	return ret;
@@ -975,15 +977,15 @@ static int edgetpu_mailbox_external_disable_free(struct edgetpu_client *client)
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (edgetpu_pm_get_if_powered(group->etdev->pm, false)) {
+	if (gcip_pm_get_if_powered(group->etdev->pm, true)) {
 		mutex_lock(&group->lock);
-		edgetpu_mailbox_external_disable_free_locked(group);
+		edgetpu_mailbox_external_free(group);
 		mutex_unlock(&group->lock);
-		edgetpu_pm_put(group->etdev->pm);
 	} else {
 		mutex_lock(&group->lock);
-		edgetpu_mailbox_external_free(group);
+		edgetpu_mailbox_external_disable_free_locked(group);
 		mutex_unlock(&group->lock);
+		gcip_pm_put(group->etdev->pm);
 	}
 
 	edgetpu_device_group_put(group);
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index 1af3bde..3478367 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -13,6 +13,8 @@
 #include <linux/of_platform.h>
 #include <linux/platform_device.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -426,7 +428,7 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 
 	dev_info(dev, "%s edgetpu initialized. Build: %s", etdev->dev_name, GIT_REPO_TAG);
 	/* Turn the device off unless a client request is already received. */
-	edgetpu_pm_shutdown(etdev, false);
+	gcip_pm_shutdown(etdev->pm, false);
 
 	edgetpu_debug_pointer = etdev;
 
@@ -445,7 +447,7 @@ out_cleanup_fw:
 	edgetpu_platform_cleanup_fw_region(etmdev);
 out_shutdown:
 	dev_dbg(dev, "Probe finished with error %d, powering down", ret);
-	edgetpu_pm_shutdown(etdev, true);
+	gcip_pm_shutdown(etdev->pm, true);
 	return ret;
 }
 
@@ -456,13 +458,13 @@ static int edgetpu_mobile_platform_remove(struct platform_device *pdev)
 
 	edgetpu_mobile_firmware_destroy(etdev);
 	edgetpu_platform_remove_irq(etmdev);
-	edgetpu_pm_get(etdev->pm);
+	gcip_pm_get(etdev->pm);
 	edgetpu_telemetry_exit(etdev);
 	edgetpu_device_remove(etdev);
 	edgetpu_iremap_pool_destroy(etdev);
 	edgetpu_platform_cleanup_fw_region(etmdev);
-	edgetpu_pm_put(etdev->pm);
-	edgetpu_pm_shutdown(etdev, true);
+	gcip_pm_put(etdev->pm);
+	gcip_pm_shutdown(etdev->pm, true);
 	edgetpu_mobile_pm_destroy(etdev);
 
 	edgetpu_debug_pointer = NULL;
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index 71cb450..bb19a42 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -5,172 +5,26 @@
  * Copyright (C) 2020 Google, Inc.
  */
 
-#include <linux/iopoll.h>
 #include <linux/mutex.h>
-#include <linux/pm.h>
-#include <linux/slab.h>
-#include <linux/workqueue.h>
 
-#include "edgetpu-config.h"
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-internal.h"
-#include "edgetpu-kci.h"
-#include "edgetpu-mailbox.h"
 #include "edgetpu-pm.h"
-#include "edgetpu-sw-watchdog.h"
 #include "edgetpu-wakelock.h"
 
-#define EDGETPU_ASYNC_POWER_DOWN_RETRY_DELAY	200
-
-struct edgetpu_pm_private {
-	const struct edgetpu_pm_handlers *handlers;
-	struct mutex lock;
-	/* Power up counter. Protected by @lock */
-	int power_up_count;
-	/* Flag indicating a deferred power down is pending. Protected by @lock */
-	bool power_down_pending;
-	/* Worker to handle async power down retry */
-	struct delayed_work power_down_work;
-	/* Back pointer to parent struct */
-	struct edgetpu_pm *etpm;
-};
-
-/*
- * Increases the counter and call the power_up callback.
- *
- * Returns zero on success.
- *
- * Caller holds etpm->p->lock.
- */
-static int edgetpu_pm_get_locked(struct edgetpu_pm *etpm)
-{
-	int power_up_count = etpm->p->power_up_count++;
-	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 (ret)
-		etpm->p->power_up_count--;
-	etdev_dbg(etpm->etdev, "%s: %d\n", __func__, etpm->p->power_up_count);
-	return ret;
-}
-
-int edgetpu_pm_trylock(struct edgetpu_pm *etpm)
-{
-	if (!etpm || !etpm->p->handlers || !etpm->p->handlers->power_up)
-		return 1;
-	return mutex_trylock(&etpm->p->lock);
-}
-
-void edgetpu_pm_unlock(struct edgetpu_pm *etpm)
-{
-	if (!etpm || !etpm->p->handlers || !etpm->p->handlers->power_up)
-		return;
-	mutex_unlock(&etpm->p->lock);
-}
-
-bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm, bool trylock)
-{
-	bool ret;
-
-	if (!etpm || !etpm->p->handlers || !etpm->p->handlers->power_up)
-		return true;
-	/* fast fail without holding the lock */
-	if (!etpm->p->power_up_count)
-		return false;
-
-	if (trylock) {
-		if (!mutex_trylock(&etpm->p->lock))
-			return false;
-	} else {
-		mutex_lock(&etpm->p->lock);
-	}
-
-	if (etpm->p->power_up_count)
-		ret = !edgetpu_pm_get_locked(etpm);
-	else
-		ret = false;
-
-	mutex_unlock(&etpm->p->lock);
-
-	return ret;
-}
-
-int edgetpu_pm_get(struct edgetpu_pm *etpm)
-{
-	int ret;
-
-	if (!etpm || !etpm->p->handlers || !etpm->p->handlers->power_up)
-		return 0;
-
-	mutex_lock(&etpm->p->lock);
-	etpm->p->power_down_pending = false;
-	ret = edgetpu_pm_get_locked(etpm);
-	mutex_unlock(&etpm->p->lock);
-
-	return ret;
-}
-
-/* Caller must hold @etpm->p->lock */
-static void edgetpu_pm_try_power_down(struct edgetpu_pm *etpm)
-{
-	int ret = etpm->p->handlers->power_down(etpm);
-
-	if (ret == -EAGAIN) {
-		etdev_warn(etpm->etdev, "Power down request denied. Retrying in %d ms\n",
-			   EDGETPU_ASYNC_POWER_DOWN_RETRY_DELAY);
-		etpm->p->power_down_pending = true;
-		schedule_delayed_work(&etpm->p->power_down_work,
-				      msecs_to_jiffies(EDGETPU_ASYNC_POWER_DOWN_RETRY_DELAY));
-	} else {
-		if (ret)
-			etdev_warn(etpm->etdev, "Power down request failed (%d)\n", ret);
-		etpm->p->power_down_pending = false;
-	}
-}
-
-/* Worker for async power down */
-static void edgetpu_pm_async_power_down_work(struct work_struct *work)
-{
-	struct delayed_work *dwork = container_of(work, struct delayed_work, work);
-	struct edgetpu_pm_private *p =
-		container_of(dwork, struct edgetpu_pm_private, power_down_work);
-
-	mutex_lock(&p->lock);
-	etdev_info(p->etpm->etdev, "Delayed power down starting\n");
-	if (p->power_down_pending)
-		edgetpu_pm_try_power_down(p->etpm);
-	else
-		etdev_info(p->etpm->etdev, "Delayed power down cancelled\n");
-	mutex_unlock(&p->lock);
-}
-
-void edgetpu_pm_put(struct edgetpu_pm *etpm)
-{
-	if (!etpm || !etpm->p->handlers || !etpm->p->handlers->power_down)
-		return;
-	mutex_lock(&etpm->p->lock);
-	if (!etpm->p->power_up_count) {
-		dev_err(etpm->etdev->dev, "Unbalanced pm_put");
-		WARN_ON(1);
-		mutex_unlock(&etpm->p->lock);
-		return;
-	}
-	if (!--etpm->p->power_up_count) {
-		edgetpu_sw_wdt_stop(etpm->etdev);
-		edgetpu_pm_try_power_down(etpm);
-	}
-	etdev_dbg(etpm->etdev, "%s: %d\n", __func__, etpm->p->power_up_count);
-	mutex_unlock(&etpm->p->lock);
-}
-
 int edgetpu_pm_create(struct edgetpu_dev *etdev,
 		      const struct edgetpu_pm_handlers *handlers)
 {
-	int ret = 0;
-	struct edgetpu_pm *etpm;
+	const struct gcip_pm_args args = {
+		.dev = etdev->dev,
+		.data = etdev,
+		.after_create = handlers->after_create,
+		.before_destroy = handlers->before_destroy,
+		.power_up = handlers->power_up,
+		.power_down = handlers->power_down,
+	};
+	struct gcip_pm *pm;
 
 	if (etdev->pm) {
 		dev_err(etdev->dev,
@@ -178,103 +32,32 @@ int edgetpu_pm_create(struct edgetpu_dev *etdev,
 		return -EEXIST;
 	}
 
-	etpm = kzalloc(sizeof(*etpm), GFP_KERNEL);
-	if (!etpm)
-		return -ENOMEM;
+	pm = gcip_pm_create(&args);
+	if (IS_ERR(pm))
+		return PTR_ERR(pm);
 
-	etpm->p = kzalloc(sizeof(*etpm->p), GFP_KERNEL);
-	if (!etpm->p) {
-		ret = -ENOMEM;
-		goto out_free_etpm;
-	}
+	etdev->pm = pm;
 
-	INIT_DELAYED_WORK(&etpm->p->power_down_work, edgetpu_pm_async_power_down_work);
-	etpm->p->etpm = etpm;
-	etpm->p->handlers = handlers;
-	etpm->etdev = etdev;
-
-	mutex_init(&etpm->p->lock);
-
-	if (handlers->after_create) {
-		ret = handlers->after_create(etpm);
-		if (ret) {
-			ret = -EINVAL;
-			goto out_free_etpm_p;
-		}
-	}
-	etdev->pm = etpm;
 	return 0;
-out_free_etpm_p:
-	kfree(etpm->p);
-out_free_etpm:
-	kfree(etpm);
-	return ret;
 }
 
 void edgetpu_pm_destroy(struct edgetpu_dev *etdev)
 {
-	const struct edgetpu_pm_handlers *handlers;
-
-	if (!etdev->pm)
-		return;
-	if (etdev->pm->p) {
-		handlers = etdev->pm->p->handlers;
-		etdev->pm->p->power_down_pending = false;
-		cancel_delayed_work_sync(&etdev->pm->p->power_down_work);
-		if (handlers && handlers->before_destroy)
-			handlers->before_destroy(etdev->pm);
-		kfree(etdev->pm->p);
-	}
-	kfree(etdev->pm);
+	gcip_pm_destroy(etdev->pm);
 	etdev->pm = NULL;
 }
 
-void edgetpu_pm_shutdown(struct edgetpu_dev *etdev, bool force)
-{
-	struct edgetpu_pm *etpm = etdev->pm;
-
-	if (!etpm)
-		return;
-
-	mutex_lock(&etpm->p->lock);
-
-	/* someone is using the device */
-	if (etpm->p->power_up_count) {
-		if (!force)
-			goto unlock;
-		else
-			etdev_warn(etdev, "Leaving %d clients behind!\n",
-				   etpm->p->power_up_count);
-	}
-
-	if (etpm->p->handlers && etpm->p->handlers->power_down)
-		etpm->p->handlers->power_down(etpm);
-unlock:
-	mutex_unlock(&etpm->p->lock);
-}
-
-bool edgetpu_is_powered(struct edgetpu_dev *etdev)
-{
-	struct edgetpu_pm *etpm = etdev->pm;
-
-	if (!etpm)
-		/* Assume powered-on in case of no power interface. */
-		return true;
-	return etpm->p->power_up_count;
-}
-
 static int __maybe_unused edgetpu_pm_suspend(struct device *dev)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_pm *etpm = etdev->pm;
+	struct gcip_pm *pm = etdev->pm;
 	struct edgetpu_list_device_client *lc;
 
-	if (!etpm || !etpm->p->power_up_count)
+	if (!pm || !gcip_pm_is_powered(pm))
 		return 0;
 
-	etdev_warn_ratelimited(
-		etdev, "cannot suspend with power up count = %d\n",
-		etpm->p->power_up_count);
+	etdev_warn_ratelimited(etdev, "cannot suspend with power up count = %d\n",
+			       gcip_pm_get_count(pm));
 
 	if (!mutex_trylock(&etdev->clients_lock))
 		return -EAGAIN;
@@ -293,14 +76,6 @@ static int __maybe_unused edgetpu_pm_suspend(struct device *dev)
 
 static int __maybe_unused edgetpu_pm_resume(struct device *dev)
 {
-	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_pm *etpm = etdev->pm;
-
-	if (etpm && etpm->p->power_up_count)
-		etdev_warn_ratelimited(etdev,
-				       "resumed with power up count = %d\n",
-				       etpm->p->power_up_count);
-
 	return 0;
 }
 
diff --git a/drivers/edgetpu/edgetpu-pm.h b/drivers/edgetpu/edgetpu-pm.h
index f0392f0..92e2ad2 100644
--- a/drivers/edgetpu/edgetpu-pm.h
+++ b/drivers/edgetpu/edgetpu-pm.h
@@ -12,66 +12,19 @@
 
 #include "edgetpu-internal.h"
 
-struct edgetpu_pm_private;
-struct edgetpu_pm;
-
 struct edgetpu_pm_handlers {
 	/* For initial setup after the interface initialized */
-	int (*after_create)(struct edgetpu_pm *etpm);
+	int (*after_create)(void *data);
 	/* For clean-up before the interface is destroyed */
-	void (*before_destroy)(struct edgetpu_pm *etpm);
+	void (*before_destroy)(void *data);
 	/* Platform-specific power up. Nesting is handled at generic layer */
-	int (*power_up)(struct edgetpu_pm *etpm);
+	int (*power_up)(void *data);
 	/* Platform-specific power down. Nesting is handled at generic layer */
-	int (*power_down)(struct edgetpu_pm *etpm);
-};
-
-struct edgetpu_pm {
-	struct edgetpu_dev *etdev;
-	struct edgetpu_pm_private *p;
+	int (*power_down)(void *data);
 };
 
 extern const struct dev_pm_ops edgetpu_pm_ops;
 
-/*
- * These mimic the pm_runtime_{get|put} functions to keep a reference count
- * of requests in order to keep the device up and turn it off if the platform
- * supports it.
- *
- * Since these functions are used by the edgetpu-fs and edgetpu-firmware
- * layers, (which have their own internal locks) no locking is provided here.
- *
- * Callers are responsible for holding any necessary locks.
- */
-
-/*
- * Tries to acquire the internal lock that ensures power_up_counter won't be
- * modified.
- *
- * Returns 1 if the lock has been acquired successfully, 0 otherwise.
- */
-int edgetpu_pm_trylock(struct edgetpu_pm *etpm);
-void edgetpu_pm_unlock(struct edgetpu_pm *etpm);
-
-/*
- * Increase power_up_count if it's already powered on.
- *
- * Caller calls edgetpu_pm_put() to decrease power_up_count if this function
- * returned true, otherwise put() shouldn't be called.
- *
- * Return false if device is not powered or trylock fails, true otherwise.
- */
-bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm, bool trylock);
-/*
- * Increase power_up_count for active state, power up the device if previous
- * power_up_count was zero.
- * Returns 0 on success or negative error value
- */
-int edgetpu_pm_get(struct edgetpu_pm *etpm);
-
-/* Decrease power_up_count for active state, power off if it reaches zero */
-void edgetpu_pm_put(struct edgetpu_pm *etpm);
-
 /* Initialize a power management interface for an edgetpu device */
 int edgetpu_pm_create(struct edgetpu_dev *etdev,
 		      const struct edgetpu_pm_handlers *handlers);
@@ -85,15 +38,4 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev);
 /* Destroy the power management interface associated with an edgetpu device */
 void edgetpu_pm_destroy(struct edgetpu_dev *etdev);
 
-/*
- * When @force is true, ensure device is shut down, regardless of whether there
- * is a client left open.
- *
- * When @force is false, the device is shut down if there is no client open.
- */
-void edgetpu_pm_shutdown(struct edgetpu_dev *etdev, bool force);
-
-/* Check if device is powered on. power_up_count is not protected by a lock */
-bool edgetpu_is_powered(struct edgetpu_dev *etdev);
-
 #endif /* __EDGETPU_PM_H__ */
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index 98bc4df..4efe0fe 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -66,7 +66,7 @@ int edgetpu_thermal_kci_if_powered(struct edgetpu_dev *etdev, u32 state);
 /*
  * Sends thermal throttling KCI to restore the last thermal state.
  *
- * The caller must guarantee the device stays powered up, typically by calling edgetpu_pm_get() or
+ * The caller must guarantee the device stays powered up, typically by calling gcip_pm_get() or
  * by calling this function from the power management functions themselves.
  *
  * Returns 0 if no thermal throttling required; otherwise the return value of KCI.
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index ab68776..bc370e5 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -14,6 +14,7 @@ gcip-objs		:= gcip-alloc-helper.o \
 			   gcip-kci.o \
 			   gcip-mailbox.o \
 			   gcip-mem-pool.o \
+			   gcip-pm.o \
 			   gcip-telemetry.o
 
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-firmware.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-firmware.c
index 0b0225c..1d9392c 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-firmware.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-firmware.c
@@ -5,7 +5,12 @@
  * Copyright (C) 2022 Google LLC
  */
 
+#include <linux/debugfs.h>
+#include <linux/mutex.h>
+#include <linux/slab.h>
+
 #include <gcip/gcip-firmware.h>
+#include <gcip/gcip-pm.h>
 
 char *gcip_fw_flavor_str(enum gcip_fw_flavor fw_flavor)
 {
@@ -23,3 +28,129 @@ char *gcip_fw_flavor_str(enum gcip_fw_flavor fw_flavor)
 		return "unknown";
 	}
 }
+
+static int gcip_firmware_tracing_active_get(void *data, u64 *val)
+{
+	struct gcip_fw_tracing *fw_tracing = data;
+
+	mutex_lock(&fw_tracing->lock);
+	*val = fw_tracing->active_level;
+	mutex_unlock(&fw_tracing->lock);
+
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_gcip_firmware_tracing_active, gcip_firmware_tracing_active_get, NULL,
+			 "%llu\n");
+
+static int gcip_firmware_tracing_request_get(void *data, u64 *val)
+{
+	struct gcip_fw_tracing *fw_tracing = data;
+
+	mutex_lock(&fw_tracing->lock);
+	*val = fw_tracing->request_level;
+	mutex_unlock(&fw_tracing->lock);
+
+	return 0;
+}
+
+static int gcip_firmware_tracing_set_level_lock(struct gcip_fw_tracing *fw_tracing)
+{
+	unsigned long active_level;
+	int ret = fw_tracing->set_level(fw_tracing->data, fw_tracing->request_level, &active_level);
+
+	if (ret)
+		dev_warn(fw_tracing->dev, "Failed to set firmware tracing level to %lu: %d",
+			 fw_tracing->request_level, ret);
+	else
+		fw_tracing->active_level =
+			(fw_tracing->request_level & GCIP_FW_TRACING_DEFAULT_VOTE) ?
+				GCIP_FW_TRACING_DEFAULT_VOTE :
+				active_level;
+
+	return ret;
+}
+
+static int gcip_firmware_tracing_request_set(void *data, u64 val)
+{
+	struct gcip_fw_tracing *fw_tracing = data;
+	int ret = 0;
+
+	mutex_lock(&fw_tracing->lock);
+
+	fw_tracing->request_level = val;
+	if (!gcip_pm_get_if_powered(fw_tracing->pm, false)) {
+		ret = gcip_firmware_tracing_set_level_lock(fw_tracing);
+		gcip_pm_put(fw_tracing->pm);
+	}
+
+	mutex_unlock(&fw_tracing->lock);
+
+	return ret;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_gcip_firmware_tracing_request, gcip_firmware_tracing_request_get,
+			 gcip_firmware_tracing_request_set, "%llu\n");
+
+struct gcip_fw_tracing *gcip_firmware_tracing_create(const struct gcip_fw_tracing_args *args)
+{
+	struct gcip_fw_tracing *fw_tracing;
+
+	if (!args->dev || !args->set_level)
+		return ERR_PTR(-EINVAL);
+
+	fw_tracing = kzalloc(sizeof(*fw_tracing), GFP_KERNEL);
+	if (!fw_tracing)
+		return ERR_PTR(-ENOMEM);
+
+	fw_tracing->dev = args->dev;
+	fw_tracing->pm = args->pm;
+	fw_tracing->set_level = args->set_level;
+	fw_tracing->data = args->data;
+	fw_tracing->active_level = GCIP_FW_TRACING_DEFAULT_VOTE;
+	fw_tracing->request_level = GCIP_FW_TRACING_DEFAULT_VOTE;
+	mutex_init(&fw_tracing->lock);
+
+	fw_tracing->dentry = debugfs_create_dir("fw_tracing", args->dentry);
+	if (IS_ERR(fw_tracing->dentry)) {
+		dev_warn(args->dev, "Failed to create debug FS tracing");
+		kfree(fw_tracing);
+
+		return (struct gcip_fw_tracing *)fw_tracing->dentry;
+	}
+
+	debugfs_create_file("active", 0440, fw_tracing->dentry, fw_tracing,
+			    &fops_gcip_firmware_tracing_active);
+	debugfs_create_file("request", 0660, fw_tracing->dentry, fw_tracing,
+			    &fops_gcip_firmware_tracing_request);
+
+	return fw_tracing;
+}
+
+void gcip_firmware_tracing_destroy(struct gcip_fw_tracing *fw_tracing)
+{
+	if (!fw_tracing)
+		return;
+
+	debugfs_remove_recursive(fw_tracing->dentry);
+	kfree(fw_tracing);
+}
+
+int gcip_firmware_tracing_restore(struct gcip_fw_tracing *fw_tracing)
+{
+	int ret = 0;
+
+	if (!fw_tracing)
+		return 0;
+
+	gcip_pm_lockdep_assert_held(fw_tracing->pm);
+	mutex_lock(&fw_tracing->lock);
+
+	fw_tracing->active_level = GCIP_FW_TRACING_DEFAULT_VOTE;
+	if (!(fw_tracing->request_level & GCIP_FW_TRACING_DEFAULT_VOTE))
+		ret = gcip_firmware_tracing_set_level_lock(fw_tracing);
+
+	mutex_unlock(&fw_tracing->lock);
+
+	return ret;
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
index cbb3c80..334a51d 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
@@ -111,10 +111,19 @@ static int gcip_mailbox_push_wait_resp(struct gcip_mailbox *mailbox, void *resp,
 {
 	struct gcip_mailbox_wait_list_elem *entry = kzalloc(sizeof(*entry), GFP_KERNEL);
 	unsigned long flags;
+	int ret;
 
 	if (!entry)
 		return -ENOMEM;
 
+	if (mailbox->ops->before_enqueue_wait_list) {
+		ret = mailbox->ops->before_enqueue_wait_list(mailbox, resp, awaiter);
+		if (ret) {
+			kfree(entry);
+			return ret;
+		}
+	}
+
 	/* Increase a reference of arrived handler. */
 	if (awaiter)
 		refcount_inc(&awaiter->refs);
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c
new file mode 100644
index 0000000..43d9654
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c
@@ -0,0 +1,239 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Power management interface for GCIP devices.
+ *
+ * Copyright (C) 2023 Google LLC
+ */
+
+#include <linux/device.h>
+#include <linux/mutex.h>
+#include <linux/slab.h>
+#include <linux/workqueue.h>
+
+#include <gcip/gcip-pm.h>
+
+#define GCIP_ASYNC_POWER_DOWN_RETRY_DELAY 200 /* ms */
+
+/* Caller must hold @pm->lock. */
+static void gcip_pm_try_power_down(struct gcip_pm *pm)
+{
+	int ret;
+
+	gcip_pm_lockdep_assert_held(pm);
+
+	ret = pm->power_down(pm->data);
+
+	if (ret == -EAGAIN) {
+		dev_warn(pm->dev, "Power down request denied, retrying in %d ms\n",
+			 GCIP_ASYNC_POWER_DOWN_RETRY_DELAY);
+		pm->power_down_pending = true;
+		schedule_delayed_work(&pm->power_down_work,
+				      msecs_to_jiffies(GCIP_ASYNC_POWER_DOWN_RETRY_DELAY));
+	} else {
+		if (ret)
+			dev_err(pm->dev, "Power down request failed (%d)\n", ret);
+		pm->power_down_pending = false;
+	}
+}
+
+/* Worker for async power down. */
+static void gcip_pm_async_power_down_work(struct work_struct *work)
+{
+	struct delayed_work *dwork = container_of(work, struct delayed_work, work);
+	struct gcip_pm *pm = container_of(dwork, struct gcip_pm, power_down_work);
+
+	mutex_lock(&pm->lock);
+
+	if (pm->power_down_pending)
+		gcip_pm_try_power_down(pm);
+	else
+		dev_info(pm->dev, "Delayed power down cancelled\n");
+
+	mutex_unlock(&pm->lock);
+}
+
+struct gcip_pm *gcip_pm_create(const struct gcip_pm_args *args)
+{
+	struct gcip_pm *pm;
+	int ret;
+
+	if (!args->dev || !args->power_up || !args->power_down)
+		return ERR_PTR(-EINVAL);
+
+	pm = devm_kzalloc(args->dev, sizeof(*pm), GFP_KERNEL);
+	if (!pm)
+		return ERR_PTR(-ENOMEM);
+
+	pm->dev = args->dev;
+	pm->data = args->data;
+	pm->after_create = args->after_create;
+	pm->before_destroy = args->before_destroy;
+	pm->power_up = args->power_up;
+	pm->power_down = args->power_down;
+
+	mutex_init(&pm->lock);
+	INIT_DELAYED_WORK(&pm->power_down_work, gcip_pm_async_power_down_work);
+
+	if (pm->after_create) {
+		ret = pm->after_create(pm->data);
+		if (ret) {
+			devm_kfree(args->dev, pm);
+			return ERR_PTR(ret);
+		}
+	}
+
+	return pm;
+}
+
+void gcip_pm_destroy(struct gcip_pm *pm)
+{
+	if (!pm)
+		return;
+
+	pm->power_down_pending = false;
+	cancel_delayed_work_sync(&pm->power_down_work);
+
+	if (pm->before_destroy)
+		pm->before_destroy(pm->data);
+
+	devm_kfree(pm->dev, pm);
+}
+
+/*
+ * Increases the counter and calls the power_up callback.
+ *
+ * Returns zero on success.
+ *
+ * Caller holds pm->lock.
+ */
+static int gcip_pm_get_locked(struct gcip_pm *pm)
+{
+	int ret = 0;
+
+	gcip_pm_lockdep_assert_held(pm);
+
+	if (!pm->count) {
+		pm->power_down_pending = false;
+		ret = pm->power_up(pm->data);
+	}
+
+	if (!ret)
+		pm->count++;
+
+	dev_dbg(pm->dev, "%s: %d\n", __func__, pm->count);
+
+	return ret;
+}
+
+int gcip_pm_get_if_powered(struct gcip_pm *pm, bool blocking)
+{
+	int ret = -EAGAIN;
+
+	if (!pm)
+		return 0;
+
+	/* Fast fails without holding the lock. */
+	if (!pm->count)
+		return ret;
+
+	if (blocking)
+		mutex_lock(&pm->lock);
+	else if (!mutex_trylock(&pm->lock))
+		return ret;
+
+	if (pm->count)
+		ret = gcip_pm_get_locked(pm);
+
+	mutex_unlock(&pm->lock);
+
+	return ret;
+}
+
+int gcip_pm_get(struct gcip_pm *pm)
+{
+	int ret;
+
+	if (!pm)
+		return 0;
+
+	mutex_lock(&pm->lock);
+	ret = gcip_pm_get_locked(pm);
+	mutex_unlock(&pm->lock);
+
+	return ret;
+}
+
+static void __gcip_pm_put(struct gcip_pm *pm, bool async)
+{
+	if (!pm)
+		return;
+
+	mutex_lock(&pm->lock);
+
+	if (WARN_ON(!pm->count))
+		goto unlock;
+
+	if (!--pm->count) {
+		pm->power_down_pending = true;
+		if (async)
+			schedule_delayed_work(&pm->power_down_work, 0);
+		else
+			gcip_pm_try_power_down(pm);
+	}
+
+	dev_dbg(pm->dev, "%s: %d\n", __func__, pm->count);
+
+unlock:
+	mutex_unlock(&pm->lock);
+}
+
+void gcip_pm_put(struct gcip_pm *pm)
+{
+	__gcip_pm_put(pm, false);
+}
+
+void gcip_pm_put_async(struct gcip_pm *pm)
+{
+	__gcip_pm_put(pm, true);
+}
+
+int gcip_pm_get_count(struct gcip_pm *pm)
+{
+	int count = -EAGAIN;
+
+	if (!pm)
+		return 0;
+
+	if (mutex_trylock(&pm->lock)) {
+		count = pm->count;
+		mutex_unlock(&pm->lock);
+	}
+
+	return count;
+}
+
+bool gcip_pm_is_powered(struct gcip_pm *pm)
+{
+	/* Assumes powered-on in case of no power interface. */
+	return pm ? gcip_pm_get_count(pm) > 0 : true;
+}
+
+void gcip_pm_shutdown(struct gcip_pm *pm, bool force)
+{
+	if (!pm)
+		return;
+
+	mutex_lock(&pm->lock);
+
+	if (pm->count) {
+		if (!force)
+			goto unlock;
+		else
+			dev_warn(pm->dev, "Force shutdown with power up count: %d", pm->count);
+	}
+
+	gcip_pm_try_power_down(pm);
+
+unlock:
+	mutex_unlock(&pm->lock);
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-firmware.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-firmware.h
index 012a79a..8cf4353 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-firmware.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-firmware.h
@@ -8,8 +8,17 @@
 #ifndef __GCIP_FIRMWARE_H__
 #define __GCIP_FIRMWARE_H__
 
+#include <linux/dcache.h>
+#include <linux/device.h>
+#include <linux/mutex.h>
 #include <linux/types.h>
 
+/*
+ * Any tracing level vote with the following bit set will be considered as a default vote.
+ * See go/gcip-firmware-dynamic-tracing for details.
+ */
+#define GCIP_FW_TRACING_DEFAULT_VOTE BIT(8)
+
 enum gcip_fw_status {
 	/* No firmware loaded yet, or last firmware failed to run. */
 	GCIP_FW_INVALID = 0,
@@ -62,4 +71,58 @@ struct gcip_fw_info {
 /* Returns the name of @fw_flavor in string. */
 char *gcip_fw_flavor_str(enum gcip_fw_flavor fw_flavor);
 
+struct gcip_fw_tracing {
+	struct device *dev;
+	struct dentry *dentry;
+	struct gcip_pm *pm;
+
+	/*
+	 * Lock to protect the struct members listed below.
+	 *
+	 * Note that since the request of tracing level adjusting might happen during power state
+	 * transitions, this lock must be acquired after holding the pm lock to avoid deadlock.
+	 */
+	struct mutex lock;
+	/* Actual firmware tracing level. */
+	unsigned long active_level;
+	/* Requested firmware tracing level. */
+	unsigned long request_level;
+
+	/* Private data. See struct gcip_fw_tracing_args.*/
+	void *data;
+
+	/* Callbacks. See struct gcip_fw_tracing_args. */
+	int (*set_level)(void *data, unsigned long level, unsigned long *active_level);
+};
+
+struct gcip_fw_tracing_args {
+	/* Device struct of GCIP device. */
+	struct device *dev;
+	/* GCIP power management. */
+	struct gcip_pm *pm;
+	/* Top-level debugfs directory for the device. */
+	struct dentry *dentry;
+	/* Private data for callbacks listed below. */
+	void *data;
+	/*
+	 * Callback to set the tracing level.
+	 * The actual tracing level clamped by the firmware should be returned by @active_level.
+	 */
+	int (*set_level)(void *data, unsigned long level, unsigned long *active_level);
+};
+
+/* Allocate and initialize the firmware tracing struct. */
+struct gcip_fw_tracing *gcip_firmware_tracing_create(const struct gcip_fw_tracing_args *args);
+
+/* Destroy and free the firmware tracing struct. */
+void gcip_firmware_tracing_destroy(struct gcip_fw_tracing *fw_tracing);
+
+/*
+ * Restore the previous firmware tracing level.
+ *
+ * This function is designed to restore the firmware tracing level during power management calls and
+ * thus it assumes the caller holds the pm lock.
+ */
+int gcip_firmware_tracing_restore(struct gcip_fw_tracing *fw_tracing);
+
 #endif /* __GCIP_FIRMWARE_H__ */
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
index e81cfb9..649b574 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
@@ -260,6 +260,22 @@ struct gcip_mailbox_ops {
 	 */
 	int (*wait_for_cmd_queue_not_full)(struct gcip_mailbox *mailbox);
 	/*
+	 * This callback will be called before putting the @resp into @mailbox->wait_list and
+	 * putting @cmd of @resp into the command queue. After this callback returns, the consumer
+	 * is able to start processing it and the mailbox is going to wait for it. Therefore, this
+	 * callback is the final checkpoint of deciding whether it is good to wait for the response
+	 * or not. If you don't want to wait for it, return a non-zero value error.
+	 *
+	 * If the implement side has its own wait queue, this callback is suitable to put @resp or
+	 * @awaiter into that.
+	 *
+	 * If @resp is synchronous, @awaiter will be NULL.
+	 *
+	 * Context: normal.
+	 */
+	int (*before_enqueue_wait_list)(struct gcip_mailbox *mailbox, void *resp,
+					struct gcip_mailbox_resp_awaiter *awaiter);
+	/*
 	 * This callback will be called after putting the @cmd to the command queue. It can be used
 	 * for triggering the doorbell. Also, @mailbox->cur_seq will be increased by the return
 	 * value. If error occurs, returns negative value and @mailbox->cur_seq will not be changed
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
new file mode 100644
index 0000000..c7673d8
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
@@ -0,0 +1,125 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Power management support for GCIP devices.
+ *
+ * Copyright (C) 2023 Google LLC
+ */
+
+#ifndef __GCIP_PM_H__
+#define __GCIP_PM_H__
+
+#include <linux/device.h>
+#include <linux/mutex.h>
+#include <linux/workqueue.h>
+
+struct gcip_pm {
+	struct device *dev;
+	/* Worker to handle async power down retry. */
+	struct delayed_work power_down_work;
+
+	/* Lock to protect the members listed below. */
+	struct mutex lock;
+	/* Power up counter. Protected by @lock */
+	int count;
+	/* Flag indicating a deferred power down is pending. Protected by @lock */
+	bool power_down_pending;
+
+	/* Callbacks. See struct gcip_pm_args. */
+	void *data;
+	int (*after_create)(void *data);
+	void (*before_destroy)(void *data);
+	int (*power_up)(void *data);
+	int (*power_down)(void *data);
+};
+
+struct gcip_pm_args {
+	/* Device struct for logging. */
+	struct device *dev;
+
+	/* Private data for the callbacks listed below. */
+	void *data;
+	/*
+	 * Device-specific power up.
+	 * Called with @pm->lock hold and nesting is handled at generic layer.
+	 */
+	int (*power_up)(void *data);
+	/*
+	 * Device-specific power down.
+	 * Called with @pm->lock hold and nesting is handled at generic layer.
+	 * Returning -EAGAIN will trigger a retry after GCIP_ASYNC_POWER_DOWN_RETRY_DELAY ms.
+	 */
+	int (*power_down)(void *data);
+	/* Optional. For initial setup after the interface initialized. */
+	int (*after_create)(void *data);
+	/* Optional. For clean-up before the interface is destroyed. */
+	void (*before_destroy)(void *data);
+};
+
+/* Allocates and initializes a power management interface for the GCIP device. */
+struct gcip_pm *gcip_pm_create(const struct gcip_pm_args *args);
+
+/* Destroys and frees the power management interface. */
+void gcip_pm_destroy(struct gcip_pm *pm);
+
+/*
+ * These mimic the pm_runtime_{get|put} functions to keep a reference count of requests in order to
+ * keep the device up and turn it off.
+ * Note that we don't keep track of system suspend/resume state since the system power management
+ * will respect the parent-child sequencing to use a bottom-up order to suspend devices and a
+ * top-down order to resume devices. No one would have the ability to acquire or release a wakelock
+ * when the device is suspending or resuming.
+ */
+
+/*
+ * Increases @pm->count if the device is already powered on.
+ *
+ * Caller should call gcip_pm_put() to decrease @pm->count if this function returns 0.
+ * If @blocking is true, it will wait until the ongoing power state transition finishes (i.e.,
+ * gcip_pm_{get,put,shutdown} called by other thread returns) and then check the power state.
+ * If @blocking is false, return -EAGAIN immediately when there is a ongoing power state transition.
+ *
+ * Returns 0 on success; otherwise -EAGAIN if the device is off or in power state transition when
+ * @blocking is false.
+ */
+int gcip_pm_get_if_powered(struct gcip_pm *pm, bool blocking);
+
+/*
+ * Increases @pm->count and powers up the device if previous @pm->count was zero.
+ *
+ * Returns 0 on success; otherwise negative error values.
+ */
+int gcip_pm_get(struct gcip_pm *pm);
+
+/*
+ * Decreases @pm->count and powers off the device if @pm->count reaches zero.
+ * If .power_down fails, async work will be scheduled to retry after
+ * GCIP_ASYNC_POWER_DOWN_RETRY_DELAY ms.
+ */
+void gcip_pm_put(struct gcip_pm *pm);
+
+/*
+ * Same as gcip_pm_put, but the power off will be scheduled later.
+ * Caller should use this async gcip_pm_put if they're on the power off path to prevent deadlock,
+ * e.g., a workqueue that will be canceled during power off.
+ */
+void gcip_pm_put_async(struct gcip_pm *pm);
+
+/* Gets the power up counter. Retures -EAGAIN if device is in power state transition. */
+int gcip_pm_get_count(struct gcip_pm *pm);
+
+/* Checks if device is already on. Retures false if device is off or in power state transition. */
+bool gcip_pm_is_powered(struct gcip_pm *pm);
+
+/* Shuts down the device if @pm->count equals to 0 or @force is true. */
+void gcip_pm_shutdown(struct gcip_pm *pm, bool force);
+
+/* Make sure @pm->lock is hold. */
+static inline void gcip_pm_lockdep_assert_held(struct gcip_pm *pm)
+{
+	if (!pm)
+		return;
+
+	lockdep_assert_held(&pm->lock);
+}
+
+#endif /* __GCIP_PM_H__ */
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 6ab65bf..0a6b194 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -18,6 +18,7 @@
 #include "edgetpu-mailbox.h"
 #include "edgetpu-mobile-platform.h"
 #include "edgetpu-pm.h"
+#include "edgetpu-sw-watchdog.h"
 #include "edgetpu-thermal.h"
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
@@ -184,11 +185,11 @@ DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, mobile_pwr_state_get, mobile_pwr_st
 DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, mobile_min_pwr_state_get, mobile_min_pwr_state_set,
 			"%llu\n");
 
-static int mobile_power_down(struct edgetpu_pm *etpm);
+static int mobile_power_down(void *data);
 
-static int mobile_power_up(struct edgetpu_pm *etpm)
+static int mobile_power_up(void *data)
 {
-	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_dev *etdev = (struct edgetpu_dev *)data;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 	int ret;
@@ -205,7 +206,7 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 			return -EAGAIN;
 	}
 
-	etdev_info(etpm->etdev, "Powering up\n");
+	etdev_info(etdev, "Powering up\n");
 
 	ret = pm_runtime_get_sync(etdev->dev);
 	if (ret) {
@@ -229,25 +230,25 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 	}
 
 	if (!etdev->firmware)
-		return 0;
+		goto out;
 
 	/*
 	 * Why this function uses edgetpu_firmware_*_locked functions without explicitly holding
 	 * edgetpu_firmware_lock:
 	 *
-	 * edgetpu_pm_get() is called in two scenarios - one is when the firmware loading is
+	 * gcip_pm_get() is called in two scenarios - one is when the firmware loading is
 	 * attempting, another one is when the user-space clients need the device be powered
 	 * (usually through acquiring the wakelock).
 	 *
 	 * For the first scenario edgetpu_firmware_is_loading() below shall return true.
 	 * For the second scenario we are indeed called without holding the firmware lock, but the
-	 * firmware loading procedures (i.e. the first scenario) always call edgetpu_pm_get() before
-	 * changing the firmware state, and edgetpu_pm_get() is blocked until this function
+	 * firmware loading procedures (i.e. the first scenario) always call gcip_pm_get() before
+	 * changing the firmware state, and gcip_pm_get() is blocked until this function
 	 * finishes. In short, we are protected by the PM lock.
 	 */
 
 	if (edgetpu_firmware_is_loading(etdev))
-		return 0;
+		goto out;
 
 	/* attempt firmware run */
 	switch (edgetpu_firmware_status_locked(etdev)) {
@@ -260,12 +261,16 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 	default:
 		break;
 	}
-	if (ret) {
-		mobile_power_down(etpm);
-	} else {
-		if (platform_pwr->post_fw_start)
-			platform_pwr->post_fw_start(etdev);
-	}
+
+	if (ret)
+		mobile_power_down(etdev);
+	else if (platform_pwr->post_fw_start)
+		platform_pwr->post_fw_start(etdev);
+
+out:
+	if (!ret)
+		edgetpu_mailbox_restore_active_mailbox_queues(etdev);
+
 	return ret;
 }
 
@@ -277,9 +282,9 @@ static void mobile_firmware_down(struct edgetpu_dev *etdev)
 		etdev_warn(etdev, "firmware shutdown failed: %d", ret);
 }
 
-static int mobile_power_down(struct edgetpu_pm *etpm)
+static int mobile_power_down(void *data)
 {
-	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_dev *etdev = (struct edgetpu_dev *)data;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 	int res = 0;
@@ -287,6 +292,8 @@ static int mobile_power_down(struct edgetpu_pm *etpm)
 
 	etdev_info(etdev, "Powering down\n");
 
+	edgetpu_sw_wdt_stop(etdev);
+
 	if (min_state >= MIN_ACTIVE_STATE) {
 		etdev_info(etdev, "Power down skipped due to min state = %d\n", min_state);
 		return 0;
@@ -342,10 +349,10 @@ static int mobile_power_down(struct edgetpu_pm *etpm)
 	return 0;
 }
 
-static int mobile_pm_after_create(struct edgetpu_pm *etpm)
+static int mobile_pm_after_create(void *data)
 {
 	int ret;
-	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_dev *etdev = (struct edgetpu_dev *)data;
 	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;
@@ -390,14 +397,14 @@ err_pm_runtime_put:
 	return ret;
 }
 
-static void mobile_pm_before_destroy(struct edgetpu_pm *etpm)
+static void mobile_pm_before_destroy(void *data)
 {
-	struct edgetpu_dev *etdev = etpm->etdev;
+	struct edgetpu_dev *etdev = (struct edgetpu_dev *)data;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 
 	debugfs_remove_recursive(platform_pwr->debugfs_dir);
-	pm_runtime_disable(etpm->etdev->dev);
+	pm_runtime_disable(etdev->dev);
 	edgetpu_soc_pm_exit(etdev);
 }
 
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index e4f1aaf..d4163ba 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -16,11 +16,12 @@
 #include <soc/google/exynos_pm_qos.h>
 #include <soc/google/gs_tmu_v3.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-internal.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mobile-platform.h"
-#include "edgetpu-pm.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-thermal.h"
 #include "mobile-firmware.h"
@@ -314,10 +315,10 @@ long edgetpu_soc_pm_get_rate(struct edgetpu_dev *etdev, int flags)
 		return -EINVAL;
 
 	/* We need to keep TPU being powered to ensure CMU read is valid. */
-	if (!edgetpu_pm_get_if_powered(etdev->pm, false))
+	if (gcip_pm_get_if_powered(etdev->pm, true))
 		return 0;
 	pll_con3 = readl(cmu_base + PLL_CON3_OFFSET);
-	edgetpu_pm_put(etdev->pm);
+	gcip_pm_put(etdev->pm);
 
 	/*
 	 * Below values must match the CMU PLL (pll_con3_pll_tpu) values in the spec and firmware.
diff --git a/drivers/edgetpu/mobile-thermal.c b/drivers/edgetpu/mobile-thermal.c
index a1bbda5..d78dc44 100644
--- a/drivers/edgetpu/mobile-thermal.c
+++ b/drivers/edgetpu/mobile-thermal.c
@@ -17,11 +17,12 @@
 #include <linux/thermal.h>
 #include <linux/version.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mmu.h"
-#include "edgetpu-pm.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-thermal.h"
 #include "mobile-pm.h"
@@ -308,13 +309,12 @@ static int edgetpu_thermal_control_kci_if_powered(struct edgetpu_dev *etdev, boo
 {
 	int ret;
 
-	/* Use trylock since this function might be called during power up. */
-	if (!edgetpu_pm_get_if_powered(etdev->pm, true))
+	if (!gcip_pm_get_if_powered(etdev->pm, false))
 		return -EAGAIN;
 
 	ret = edgetpu_thermal_control_kci(etdev, enable);
 
-	edgetpu_pm_put(etdev->pm);
+	gcip_pm_put(etdev->pm);
 
 	return ret;
 }
@@ -497,12 +497,11 @@ int edgetpu_thermal_kci_if_powered(struct edgetpu_dev *etdev, u32 state)
 {
 	int ret;
 
-	/* Use trylock since this function might be called during power up. */
-	if (!edgetpu_pm_get_if_powered(etdev->pm, true))
+	if (gcip_pm_get_if_powered(etdev->pm, false))
 		return -EAGAIN;
 
 	ret = edgetpu_thermal_kci(etdev, state);
-	edgetpu_pm_put(etdev->pm);
+	gcip_pm_put(etdev->pm);
 
 	return ret;
 }
diff --git a/drivers/edgetpu/rio-pm.c b/drivers/edgetpu/rio-pm.c
index 58271d7..c013583 100644
--- a/drivers/edgetpu/rio-pm.c
+++ b/drivers/edgetpu/rio-pm.c
@@ -12,6 +12,7 @@
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-mobile-platform.h"
+#include "edgetpu-pm.h"
 #include "edgetpu-soc.h"
 #include "mobile-soc-gsx01.h"
 #include "mobile-pm.h"