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

gxp: remove redundant domain_attach from vd resume
Bug: 269587251
gcip: fixup: Use devm_* for gcip-pm
Bug: 265870718
gcip: Use devm_* for gcip-pm
Bug: 265870718 (repeat)
gcip: Remove the start log of async power down
Bug: 265870718 (repeat)
gcip: Add firmware dynamic tracing support
Bug: 262916889
gcip: Add PM support
Bug: 265870718 (repeat)
GCIP_MAIN_REV_ID: c359c8b4c8e11ff2655dbfd8457605b760db383c
gxp: bump version 1.10
Bug: 269587251 (repeat)
gxp: move domain detach to block_unready
Bug: 269587251 (repeat)
gxp: add mcu_crashed flag to gxp_virtual_device
Bug: 269587251 (repeat)
gxp: use MCU PSM to handle MCU fw crash
Bug: 264621513
gxp: add IDs to VD suspend/resume log
gxp: skip suspend/resume in full MCU mode
Bug: 269717931
gxp: load core FW only once per image
Bug: 267713927
gxp: Adopt GCIP PM
Bug: 265870718 (repeat)
gxp: Use gxp_wakelock_* in mcu fw crash handler
Bug: 265870718 (repeat)
gxp: Log the clients with a wakelock in gxp_platform_suspend
Bug: 265870718 (repeat)
gxp: Remove suspended from gxp_wakelock_manager
Bug: 265870718 (repeat)
gxp: support loading DSP core FW in raw binary format
Bug: 259215977
gxp: release_vmbox handles ABORTED error
Bug: 263994153
Bug: 267697081
gxp: introduce gxp_vd_invalidate_with_client_id
Bug: 263994153 (repeat)
gxp: introduce gxp_vd_generate_debug_dump
Bug: 263994153 (repeat)
gxp: move vd_invalid_eventfd from client to vd
Bug: 263994153 (repeat)
gxp: gxp_vd_block_unready needs vd_semaphore write lock
Bug: 263994153 (repeat)
gcip: Add a comment about suspend/resume in gcip-pm
Bug: 265870718 (repeat)
gcip: Add firmware dynamic tracing header
Bug: 262916889 (repeat)
gcip: Add PM header
Bug: 265870718 (repeat)
GCIP_HEADERS_REV_ID: 21105099cc4be7a0fec32f06f34cce5026015601
GitOrigin-RevId: 33992f22cfe1d039eec901aa6b937d6e0e5d4c3d
Change-Id: I2ca86fd3df6415641e87a830c7915288f0e50de3

26 files changed, 948 insertions, 482 deletions

diff --git a/Makefile b/Makefile
index c1b4ca2..472c13e 100644
--- a/Makefile
+++ b/Makefile
@@ -32,8 +32,7 @@ gxp-objs += \
 		gxp-range-alloc.o \
 		gxp-ssmt.o \
 		gxp-thermal.o \
-		gxp-vd.o \
-		gxp-wakelock.o
+		gxp-vd.o
 
 
 ifeq ($(GXP_CHIP),CALLISTO)
diff --git a/callisto/config.h b/callisto/config.h
index bcae5fd..6b9aa0d 100644
--- a/callisto/config.h
+++ b/callisto/config.h
@@ -37,8 +37,6 @@
 /* The max number of active virtual devices. */
 #define GXP_NUM_SHARED_SLICES 7
 
-/* TODO(b/234098135): remove this when FW supports suspend / resume */
-#define DISABLE_VD_SUSPEND_RESUME_SUPPORT
 /*
  * Can be coherent with AP
  *
diff --git a/gcip-kernel-driver/drivers/gcip/Makefile b/gcip-kernel-driver/drivers/gcip/Makefile
index ab68776..bc370e5 100644
--- a/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/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/gcip-kernel-driver/drivers/gcip/gcip-firmware.c b/gcip-kernel-driver/drivers/gcip/gcip-firmware.c
index 0b0225c..1d9392c 100644
--- a/gcip-kernel-driver/drivers/gcip/gcip-firmware.c
+++ b/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/gcip-kernel-driver/drivers/gcip/gcip-pm.c b/gcip-kernel-driver/drivers/gcip/gcip-pm.c
new file mode 100644
index 0000000..43d9654
--- /dev/null
+++ b/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/gcip-kernel-driver/include/gcip/gcip-firmware.h b/gcip-kernel-driver/include/gcip/gcip-firmware.h
index 012a79a..8cf4353 100644
--- a/gcip-kernel-driver/include/gcip/gcip-firmware.h
+++ b/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/gcip-kernel-driver/include/gcip/gcip-pm.h b/gcip-kernel-driver/include/gcip/gcip-pm.h
new file mode 100644
index 0000000..c7673d8
--- /dev/null
+++ b/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/gxp-client.c b/gxp-client.c
index 9686818..7c96e9a 100644
--- a/gxp-client.c
+++ b/gxp-client.c
@@ -9,12 +9,14 @@
 #include <linux/slab.h>
 #include <linux/types.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "gxp-client.h"
 #include "gxp-dma.h"
 #include "gxp-internal.h"
 #include "gxp-pm.h"
 #include "gxp-vd.h"
-#include "gxp-wakelock.h"
+#include "gxp.h"
 
 struct gxp_client *gxp_client_create(struct gxp_dev *gxp)
 {
@@ -40,23 +42,23 @@ void gxp_client_destroy(struct gxp_client *client)
 	struct gxp_dev *gxp = client->gxp;
 	int core;
 
-	if (client->vd && client->has_block_wakelock)
-		gxp_vd_block_unready(client->vd);
-
 	if (client->vd && client->vd->state != GXP_VD_OFF) {
 		down_write(&gxp->vd_semaphore);
 		gxp_vd_stop(client->vd);
 		up_write(&gxp->vd_semaphore);
 	}
 
+	if (client->vd && client->has_block_wakelock) {
+		down_write(&gxp->vd_semaphore);
+		gxp_vd_block_unready(client->vd);
+		up_write(&gxp->vd_semaphore);
+	}
+
 	for (core = 0; core < GXP_NUM_CORES; core++) {
 		if (client->mb_eventfds[core])
 			gxp_eventfd_put(client->mb_eventfds[core]);
 	}
 
-	if (client->vd_invalid_eventfd)
-		gxp_eventfd_put(client->vd_invalid_eventfd);
-
 #if (IS_ENABLED(CONFIG_GXP_TEST) || IS_ENABLED(CONFIG_ANDROID)) &&             \
 	!IS_ENABLED(CONFIG_GXP_GEM5)
 	if (client->tpu_file) {
@@ -77,7 +79,7 @@ void gxp_client_destroy(struct gxp_client *client)
 #endif
 
 	if (client->has_block_wakelock) {
-		gxp_wakelock_release(client->gxp);
+		gcip_pm_put(client->gxp->power_mgr->pm);
 		gxp_pm_update_requested_power_states(
 			gxp, client->requested_states, off_states);
 	}
@@ -181,7 +183,7 @@ int gxp_client_acquire_block_wakelock(struct gxp_client *client,
 
 	lockdep_assert_held(&client->semaphore);
 	if (!client->has_block_wakelock) {
-		ret = gxp_wakelock_acquire(gxp);
+		ret = gcip_pm_get(gxp->power_mgr->pm);
 		if (ret)
 			return ret;
 		*acquired_wakelock = true;
@@ -208,7 +210,7 @@ int gxp_client_acquire_block_wakelock(struct gxp_client *client,
 
 err_wakelock_release:
 	if (*acquired_wakelock) {
-		gxp_wakelock_release(gxp);
+		gcip_pm_put(gxp->power_mgr->pm);
 		*acquired_wakelock = false;
 	}
 	return ret;
@@ -222,13 +224,15 @@ void gxp_client_release_block_wakelock(struct gxp_client *client)
 	if (!client->has_block_wakelock)
 		return;
 
-	if (client->vd)
-		gxp_vd_block_unready(client->vd);
+	gxp_client_release_vd_wakelock(client);
 
-	if (client->has_vd_wakelock)
-		gxp_client_release_vd_wakelock(client);
+	if (client->vd) {
+		down_write(&gxp->vd_semaphore);
+		gxp_vd_block_unready(client->vd);
+		up_write(&gxp->vd_semaphore);
+	}
 
-	gxp_wakelock_release(gxp);
+	gcip_pm_put(gxp->power_mgr->pm);
 	client->has_block_wakelock = false;
 }
 
diff --git a/gxp-client.h b/gxp-client.h
index 935b23d..964aa76 100644
--- a/gxp-client.h
+++ b/gxp-client.h
@@ -40,7 +40,6 @@ struct gxp_client {
 	struct gxp_tpu_mbx_desc mbx_desc;
 
 	struct gxp_eventfd *mb_eventfds[GXP_NUM_CORES];
-	struct gxp_eventfd *vd_invalid_eventfd;
 
 	/* client process thread group ID is really the main process ID. */
 	pid_t tgid;
diff --git a/gxp-common-platform.c b/gxp-common-platform.c
index 68accfa..de3417c 100644
--- a/gxp-common-platform.c
+++ b/gxp-common-platform.c
@@ -23,6 +23,7 @@
 #include <linux/uidgid.h>
 
 #include <gcip/gcip-dma-fence.h>
+#include <gcip/gcip-pm.h>
 
 #include "gxp-client.h"
 #include "gxp-config.h"
@@ -44,7 +45,6 @@
 #include "gxp-pm.h"
 #include "gxp-thermal.h"
 #include "gxp-vd.h"
-#include "gxp-wakelock.h"
 #include "gxp.h"
 
 #if HAS_TPU_EXT
@@ -1614,16 +1614,21 @@ static int gxp_register_invalidated_eventfd(
 
 	down_write(&client->semaphore);
 
+	if (!gxp_client_has_available_vd(client,
+					 "GXP_REGISTER_INVALIDATED_EVENTFD")) {
+		ret = -ENODEV;
+		goto out;
+	}
+
 	eventfd = gxp_eventfd_create(ibuf.eventfd);
 	if (IS_ERR(eventfd)) {
 		ret = PTR_ERR(eventfd);
 		goto out;
 	}
 
-	if (client->vd_invalid_eventfd)
-		gxp_eventfd_put(client->vd_invalid_eventfd);
-	client->vd_invalid_eventfd = eventfd;
-
+	if (client->vd->invalidate_eventfd)
+		gxp_eventfd_put(client->vd->invalidate_eventfd);
+	client->vd->invalidate_eventfd = eventfd;
 out:
 	up_write(&client->semaphore);
 	return ret;
@@ -1633,14 +1638,24 @@ static int gxp_unregister_invalidated_eventfd(
 	struct gxp_client *client,
 	struct gxp_register_invalidated_eventfd_ioctl __user *argp)
 {
+	struct gxp_dev *gxp = client->gxp;
+	int ret = 0;
+
 	down_write(&client->semaphore);
 
-	if (client->vd_invalid_eventfd)
-		gxp_eventfd_put(client->vd_invalid_eventfd);
-	client->vd_invalid_eventfd = NULL;
+	if (!client->vd) {
+		dev_err(gxp->dev,
+			"GXP_UNREGISTER_INVALIDATED_EVENTFD requires the client allocate a VIRTUAL_DEVICE\n");
+		ret = -ENODEV;
+		goto out;
+	}
 
+	if (client->vd->invalidate_eventfd)
+		gxp_eventfd_put(client->vd->invalidate_eventfd);
+	client->vd->invalidate_eventfd = NULL;
+out:
 	up_write(&client->semaphore);
-	return 0;
+	return ret;
 }
 
 static long gxp_ioctl(struct file *file, uint cmd, ulong arg)
@@ -1980,16 +1995,11 @@ static int gxp_common_platform_probe(struct platform_device *pdev, struct gxp_de
 		return ret;
 
 	gxp_create_debugdir(gxp);
-	ret = gxp_wakelock_init(gxp);
-	if (ret) {
-		dev_err(dev, "failed to init wakelock: %d", ret);
-		goto err_remove_debugdir;
-	}
 
 	ret = gxp_pm_init(gxp);
 	if (ret) {
 		dev_err(dev, "Failed to init power management (ret=%d)\n", ret);
-		goto err_wakelock_destroy;
+		goto err_remove_debugdir;
 	}
 
 	gxp_get_gsa_dev(gxp);
@@ -2137,8 +2147,6 @@ err_put_tpu_dev:
 	gxp_put_tpu_dev(gxp);
 	gxp_put_gsa_dev(gxp);
 	gxp_pm_destroy(gxp);
-err_wakelock_destroy:
-	/* wakelock init doesn't need revert */
 err_remove_debugdir:
 	gxp_remove_debugdir(gxp);
 	return ret;
@@ -2174,15 +2182,45 @@ static int gxp_common_platform_remove(struct platform_device *pdev)
 static int gxp_platform_suspend(struct device *dev)
 {
 	struct gxp_dev *gxp = dev_get_drvdata(dev);
+	struct gxp_client *client;
+
+	if (!gcip_pm_is_powered(gxp->power_mgr->pm))
+		return 0;
+
+	/* Log clients currently holding a wakelock */
+	if (!mutex_trylock(&gxp->client_list_lock)) {
+		dev_warn_ratelimited(
+			gxp->dev,
+			"Unable to get client list lock on suspend failure\n");
+		return -EAGAIN;
+	}
+
+	list_for_each_entry(client, &gxp->client_list, list_entry) {
+		if (!down_read_trylock(&client->semaphore)) {
+			dev_warn_ratelimited(
+				gxp->dev,
+				"Unable to acquire client lock (tgid=%d pid=%d)\n",
+				client->tgid, client->pid);
+			continue;
+		}
+
+		if (client->has_block_wakelock)
+			dev_warn_ratelimited(
+				gxp->dev,
+				"Cannot suspend with client holding wakelock (tgid=%d pid=%d)\n",
+				client->tgid, client->pid);
 
-	return gxp_wakelock_suspend(gxp);
+		up_read(&client->semaphore);
+	}
+
+	mutex_unlock(&gxp->client_list_lock);
+
+	return -EAGAIN;
 }
 
 static int gxp_platform_resume(struct device *dev)
 {
-	struct gxp_dev *gxp = dev_get_drvdata(dev);
-
-	return gxp_wakelock_resume(gxp);
+	return 0;
 }
 
 static const struct dev_pm_ops gxp_pm_ops = {
diff --git a/gxp-debug-dump.c b/gxp-debug-dump.c
index 4c2483b..94ae78b 100644
--- a/gxp-debug-dump.c
+++ b/gxp-debug-dump.c
@@ -18,6 +18,8 @@
 #include <linux/platform_data/sscoredump.h>
 #endif
 
+#include <gcip/gcip-pm.h>
+
 #include "gxp-client.h"
 #include "gxp-debug-dump.h"
 #include "gxp-dma.h"
@@ -29,7 +31,6 @@
 #include "gxp-mapping.h"
 #include "gxp-pm.h"
 #include "gxp-vd.h"
-#include "gxp-wakelock.h"
 
 #define SSCD_MSG_LENGTH 64
 
@@ -273,7 +274,7 @@ static int gxp_get_common_dump(struct gxp_dev *gxp)
 	int ret;
 
 	/* Power on BLK_AUR to read the common registers */
-	ret = gxp_wakelock_acquire(gxp);
+	ret = gcip_pm_get(gxp->power_mgr->pm);
 	if (ret) {
 		dev_err(gxp->dev,
 			"Failed to acquire wakelock for getting common dump\n");
@@ -287,7 +288,7 @@ static int gxp_get_common_dump(struct gxp_dev *gxp)
 	gxp_get_lpm_registers(gxp, &common_seg_header[GXP_LPM_REGISTERS_IDX],
 			      &common_dump_data->lpm_regs);
 
-	gxp_wakelock_release(gxp);
+	gcip_pm_put(gxp->power_mgr->pm);
 	gxp_pm_update_requested_power_states(gxp, uud_states, off_states);
 
 	dev_dbg(gxp->dev, "Segment Header for Common Segment\n");
diff --git a/gxp-debugfs.c b/gxp-debugfs.c
index fa1685e..a9a2e14 100644
--- a/gxp-debugfs.c
+++ b/gxp-debugfs.c
@@ -7,6 +7,8 @@
 
 #include <linux/acpm_dvfs.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "gxp-client.h"
 #include "gxp-core-telemetry.h"
 #include "gxp-debug-dump.h"
@@ -20,7 +22,6 @@
 #include "gxp-mailbox.h"
 #include "gxp-pm.h"
 #include "gxp-vd.h"
-#include "gxp-wakelock.h"
 #include "gxp.h"
 
 #if GXP_HAS_MCU
@@ -266,11 +267,9 @@ static int gxp_wakelock_set(void *data, u64 val)
 			goto out;
 		}
 
-		ret = gxp_wakelock_acquire(gxp);
+		ret = gcip_pm_get(gxp->power_mgr->pm);
 		if (ret) {
-			dev_err(gxp->dev,
-				"Failed to acquire debugfs wakelock ret=%d\n",
-				ret);
+			dev_err(gxp->dev, "gcip_pm_get failed ret=%d\n", ret);
 			goto out;
 		}
 		gxp->debugfs_wakelock_held = true;
@@ -284,7 +283,7 @@ static int gxp_wakelock_set(void *data, u64 val)
 			goto out;
 		}
 
-		gxp_wakelock_release(gxp);
+		gcip_pm_put(gxp->power_mgr->pm);
 		gxp->debugfs_wakelock_held = false;
 		gxp_pm_update_requested_power_states(gxp, uud_states,
 						     off_states);
diff --git a/gxp-firmware.c b/gxp-firmware.c
index 4a59a13..b0453d5 100644
--- a/gxp-firmware.c
+++ b/gxp-firmware.c
@@ -138,8 +138,8 @@ static int elf_load_segments(struct gxp_dev *gxp, const u8 *elf_data,
 	    (ehdr->e_ident[EI_MAG1] != ELFMAG1) ||
 	    (ehdr->e_ident[EI_MAG2] != ELFMAG2) ||
 	    (ehdr->e_ident[EI_MAG3] != ELFMAG3)) {
-		dev_err(gxp->dev, "Invalid ELF format.");
-		return -EINVAL;
+		dev_info(gxp->dev, "Firmware is not an ELF, treated as raw binary.");
+		return 0;
 	}
 
 	/* go through the available ELF segments */
@@ -213,15 +213,6 @@ static int elf_load_segments(struct gxp_dev *gxp, const u8 *elf_data,
 	return ret;
 }
 
-static void elf_fetch_entry_point(struct gxp_dev *gxp, const u8 *elf_data,
-				  uint core)
-{
-	struct elf32_hdr *ehdr;
-
-	ehdr = (struct elf32_hdr *)elf_data;
-	gxp->firmware_mgr->entry_points[core] = ehdr->e_entry;
-}
-
 static int
 gxp_firmware_authenticate(struct gxp_dev *gxp,
 			  const struct firmware *firmwares[GXP_NUM_CORES])
@@ -322,14 +313,14 @@ static void gxp_program_reset_vector(struct gxp_dev *gxp, uint core,
 	reset_vec = gxp_read_32(gxp, GXP_CORE_REG_ALT_RESET_VECTOR(phys_core));
 	if (verbose)
 		dev_notice(gxp->dev,
-			   "Current Aurora reset vector for core %u: 0x%x\n",
+			   "Current Aurora reset vector for core %u: %#x\n",
 			   phys_core, reset_vec);
 	gxp_write_32(gxp, GXP_CORE_REG_ALT_RESET_VECTOR(phys_core),
-		     gxp->firmware_mgr->entry_points[core]);
+		     gxp->fwbufs[core].daddr);
 	if (verbose)
 		dev_notice(gxp->dev,
-			   "New Aurora reset vector for core %u: 0x%x\n",
-			   phys_core, gxp->firmware_mgr->entry_points[core]);
+			   "New Aurora reset vector for core %u: %#llx\n",
+			   phys_core, gxp->fwbufs[core].daddr);
 }
 
 static void *get_scratchpad_base(struct gxp_dev *gxp,
@@ -400,6 +391,8 @@ static int gxp_firmware_load(struct gxp_dev *gxp, struct gxp_virtual_device *vd,
 
 	if (!mgr->firmwares[core])
 		return -ENODEV;
+	if (mgr->loaded[core])
+		return 0;
 
 	/* Load firmware to System RAM */
 	ret = elf_load_segments(gxp,
@@ -411,14 +404,11 @@ static int gxp_firmware_load(struct gxp_dev *gxp, struct gxp_virtual_device *vd,
 		goto out_firmware_unload;
 	}
 
-	elf_fetch_entry_point(gxp,
-			      mgr->firmwares[core]->data + FW_HEADER_SIZE,
-			      core);
-
 	/* TODO(b/188970444): Cleanup logging of addresses */
 	dev_notice(gxp->dev,
-		   "ELF loaded at virtual: %pK and physical: 0x%llx\n",
+		   "ELF loaded at virtual: %pK and physical: %#llx\n",
 		   gxp->fwbufs[core].vaddr, gxp->fwbufs[core].paddr);
+	mgr->loaded[core] = true;
 
 	return 0;
 
@@ -619,6 +609,7 @@ static ssize_t load_dsp_firmware_store(struct device *dev,
 		if (mgr->firmwares[core])
 			release_firmware(mgr->firmwares[core]);
 		mgr->firmwares[core] = firmwares[core];
+		mgr->loaded[core] = false;
 	}
 
 	kfree(mgr->firmware_name);
diff --git a/gxp-firmware.h b/gxp-firmware.h
index f19adc9..fdaff99 100644
--- a/gxp-firmware.h
+++ b/gxp-firmware.h
@@ -58,10 +58,10 @@ struct gxp_firmware_manager {
 	bool is_firmware_requested;
 	/* Protects `firmwares` and `firmware_name` */
 	struct mutex dsp_firmware_lock;
+	/* FW is readonly, we only need to load it once per image. */
+	bool loaded[GXP_NUM_CORES];
 	/* Firmware status bitmap. Accessors must hold `vd_semaphore`. */
 	u32 firmware_running;
-	/* Store the entry point of the DSP core firmware. */
-	u32 entry_points[GXP_NUM_CORES];
 	/*
 	 * Cached image config, for easier fetching config entries.
 	 * Not a pointer to the firmware buffer because we want to forcely change the
diff --git a/gxp-internal.h b/gxp-internal.h
index 1eaa5c0..cc6cffe 100644
--- a/gxp-internal.h
+++ b/gxp-internal.h
@@ -66,7 +66,6 @@ struct gxp_fw_data_manager;
 struct gxp_power_manager;
 struct gxp_core_telemetry_manager;
 struct gxp_thermal_manager;
-struct gxp_wakelock_manager;
 struct gxp_usage_stats;
 struct gxp_power_states;
 struct gxp_iommu_domain;
@@ -110,7 +109,6 @@ struct gxp_dev {
 	struct gxp_fw_data_manager *data_mgr;
 	struct gxp_tpu_dev tpu_dev;
 	struct gxp_core_telemetry_manager *core_telemetry_mgr;
-	struct gxp_wakelock_manager *wakelock_mgr;
 	struct gxp_iommu_domain *default_domain;
 	/*
 	 * Pointer to GSA device for firmware authentication.
@@ -199,6 +197,7 @@ struct gxp_dev {
 				    struct gxp_power_states power_states);
 	/*
 	 * Called when the client acquired the BLOCK wakelock and allocated a virtual device.
+	 * The caller will hold @gxp->vd_semaphore for writing.
 	 *
 	 * Return a non-zero value can fail the block acquiring.
 	 *
@@ -208,29 +207,30 @@ struct gxp_dev {
 				    struct gxp_virtual_device *vd);
 	/*
 	 * Called before releasing the BLOCK wakelock or the virtual device.
+	 * The caller will hold @gxp->vd_semaphore for writing.
 	 *
-	 * This callback is optional
+	 * This callback is optional.
 	 */
 	void (*before_vd_block_unready)(struct gxp_dev *gxp,
 					struct gxp_virtual_device *vd);
 	/*
-	 * Called in gxp_wakelock_acquire(), after the block is powered.
+	 * Called in .power_up callback of gcip_pm, after the block is powered.
 	 *
-	 * This function is called with holding gxp_wakelock_manager.lock.
+	 * This function is called with holding gcip_pm lock.
 	 *
-	 * Return a non-zero value can fail gxp_wakelock_acquire().
+	 * Return a non-zero value can fail gcip_pm_get.
 	 *
 	 * This callback is optional.
 	 */
-	int (*wakelock_after_blk_on)(struct gxp_dev *gxp);
+	int (*pm_after_blk_on)(struct gxp_dev *gxp);
 	/*
-	 * Called in gxp_wakelock_release(), before the block is shutdown.
+	 * Called in .power_down callback of gcip_pm, before the block is shutdown.
 	 *
-	 * This function is called with holding gxp_wakelock_manager.lock.
+	 * This function is called with holding gcip_pm lock.
 	 *
 	 * This callback is optional.
 	 */
-	void (*wakelock_before_blk_off)(struct gxp_dev *gxp);
+	void (*pm_before_blk_off)(struct gxp_dev *gxp);
 	/*
 	 * Called in gxp_map_tpu_mbx_queue(), after the TPU mailbox buffers are mapped.
 	 *
diff --git a/gxp-kci.c b/gxp-kci.c
index 5ff365b..6e19686 100644
--- a/gxp-kci.c
+++ b/gxp-kci.c
@@ -142,15 +142,15 @@ static void gxp_kci_handle_rkci(struct gxp_kci *gkci,
 		if (client_id == SECURE_CLIENT_ID)
 			core_list = 0;
 		/*
-		 * Inside gxp_vd_invalidate() after invalidating the client, debug dump
-		 * if enabled would be checked and processed for individual cores in
+		 * Inside gxp_vd_invalidate_with_client_id() after invalidating the client, debug
+		 * dump if enabled would be checked and processed for individual cores in
 		 * core_list. Due to debug dump processing being a time consuming task
 		 * rkci ack is sent first to unblock the mcu to send furhter rkci's. Client
-		 * lock inside gxp_vd_invalidate() would make sure the correctness of the
-		 * logic against possible concurrent scenarios.
+		 * lock inside gxp_vd_invalidate_with_client_id() would make sure the correctness
+		 * of the logic against possible concurrent scenarios.
 		 */
 		gxp_kci_resp_rkci_ack(gkci, resp);
-		gxp_vd_invalidate(gxp, client_id, core_list);
+		gxp_vd_invalidate_with_client_id(gxp, client_id, core_list);
 
 		break;
 	}
diff --git a/gxp-mcu-firmware.c b/gxp-mcu-firmware.c
index be0fbad..6e1031a 100644
--- a/gxp-mcu-firmware.c
+++ b/gxp-mcu-firmware.c
@@ -15,6 +15,7 @@
 
 #include <gcip/gcip-common-image-header.h>
 #include <gcip/gcip-image-config.h>
+#include <gcip/gcip-pm.h>
 
 #include "gxp-bpm.h"
 #include "gxp-config.h"
@@ -26,7 +27,6 @@
 #include "gxp-mcu-firmware.h"
 #include "gxp-mcu.h"
 #include "gxp-pm.h"
-#include "gxp-wakelock.h"
 
 /* Value of Magic field in the common header "DSPF' as a 32-bit LE int */
 #define GXP_FW_MAGIC 0x46505344
@@ -377,7 +377,7 @@ static ssize_t load_firmware_store(struct device *dev,
 	name = fw_name_from_buf(gxp, buf);
 	if (IS_ERR(name))
 		return PTR_ERR(name);
-	if (gxp->wakelock_mgr->count) {
+	if (gcip_pm_is_powered(gxp->power_mgr->pm)) {
 		dev_err(gxp->dev,
 			"Reject firmware loading because wakelocks are holding");
 		return -EBUSY;
@@ -391,12 +391,12 @@ static ssize_t load_firmware_store(struct device *dev,
 	dev_info(gxp->dev, "loading firmware %s from SysFS", name);
 	last_name = mcu_fw->name;
 	mcu_fw->name = name;
-	ret = gxp_wakelock_acquire(gxp);
+	ret = gcip_pm_get(gxp->power_mgr->pm);
 	if (ret) {
 		dev_err(gxp->dev, "loading firmware %s failed: %d", name, ret);
 		mcu_fw->name = last_name;
 	} else {
-		gxp_wakelock_release(gxp);
+		gcip_pm_put(gxp->power_mgr->pm);
 	}
 	return ret < 0 ? ret : count;
 }
@@ -566,23 +566,20 @@ void gxp_mcu_firmware_crash_handler(struct gxp_dev *gxp,
 	down_write(&gxp->vd_semaphore);
 
 	/*
-	 * As we are recovering the MCU firmware, the number of clients holding the block wakelock
-	 * should not be changed until the rescuing is finished.
+	 * As we are recovering the MCU firmware, we should prevent power state transition caused by
+	 * clients acquiring or releasing the block wakelock until the rescuing is finished.
+	 *
+	 * The runtime can still acquire or release the block wakelock, but commands issued to the
+	 * firmware might not be handled properly.
 	 *
-	 * The runtime cannot acquire or release the block wakelock until this function releases
-	 * the lock.
-	 */
-	mutex_lock(&gxp->wakelock_mgr->lock);
-
-	/*
 	 * By the race, if all clients left earlier than this handler, all block wakleock should be
 	 * already released and the BLK is turned off. We don't have to rescue the MCU firmware.
 	 */
-	if (!gxp->wakelock_mgr->count) {
+	if (gcip_pm_get_if_powered(gxp->power_mgr->pm, false)) {
 		dev_info(
 			gxp->dev,
 			"The block wakelock is already released, skip restarting MCU firmware");
-		goto unlock_wakelock_mgr;
+		goto out_unlock_client_semaphore;
 	}
 
 	/*
@@ -591,36 +588,25 @@ void gxp_mcu_firmware_crash_handler(struct gxp_dev *gxp,
 	 */
 	list_for_each_entry (client, &gxp->client_list, list_entry) {
 		if (client->has_block_wakelock && client->vd) {
-			gxp->mailbox_mgr->release_unconsumed_async_resps(
-				client->vd);
-			client->vd->state = GXP_VD_UNAVAILABLE;
-			if (client->vd_invalid_eventfd)
-				gxp_eventfd_signal(client->vd_invalid_eventfd);
+			gxp_vd_invalidate(gxp, client->vd);
+			client->vd->mcu_crashed = true;
 		}
 	}
 
-	/*
-	 * Turn off and on the Aurora block and rerun the MCU firmware.
-	 * TODO(b/264621513): Change the power state of LPM instead of turning off and on the
-	 * whole Aurora block.
-	 */
+	/* Turn off and on the MCU PSM and restart the MCU firmware. */
 	mutex_lock(&mcu_fw->lock);
 
-	ret = gxp_pm_blk_off(gxp);
-	if (ret) {
-		dev_err(gxp->dev, "Failed to turn off BLK_AUR (ret=%d)\n", ret);
-		goto out;
-	}
-
-	if (!gxp_pm_is_blk_down(gxp, 5000)) {
-		dev_err(gxp->dev, "BLK_AUR hasn't been turned off");
-		goto out;
-	}
+	gxp_lpm_down(gxp, GXP_MCU_CORE_ID);
 
-	ret = gxp_pm_blk_on(gxp);
-	if (ret) {
-		dev_err(gxp->dev, "Failed to turn on BLK_AUR (ret=%d)\n", ret);
-		goto out;
+	if (!gxp_lpm_wait_state_eq(gxp, CORE_TO_PSM(GXP_MCU_CORE_ID),
+				   LPM_PG_STATE)) {
+		dev_warn(
+			gxp->dev,
+			"MCU PSM transition to PS3 fails, current state: %u. Falling back to power cycle AUR block.\n",
+			gxp_lpm_get_state(gxp, CORE_TO_PSM(GXP_MCU_CORE_ID)));
+		ret = gxp_pm_blk_reboot(gxp, 5000);
+		if (ret)
+			goto out;
 	}
 
 	ret = gxp_mcu_firmware_restart_locked(mcu_fw);
@@ -629,8 +615,8 @@ void gxp_mcu_firmware_crash_handler(struct gxp_dev *gxp,
 
 out:
 	mutex_unlock(&mcu_fw->lock);
-unlock_wakelock_mgr:
-	mutex_unlock(&gxp->wakelock_mgr->lock);
+	gcip_pm_put(gxp->power_mgr->pm);
+out_unlock_client_semaphore:
 	up_write(&gxp->vd_semaphore);
 	list_for_each_entry (client, &gxp->client_list, list_entry) {
 		up_write(&client->semaphore);
diff --git a/gxp-mcu-platform.c b/gxp-mcu-platform.c
index de16dc4..3a7c7f4 100644
--- a/gxp-mcu-platform.c
+++ b/gxp-mcu-platform.c
@@ -14,6 +14,13 @@
 #include "gxp-mcu.h"
 #include "gxp-usage-stats.h"
 
+#define KCI_RETURN_CORE_LIST_MASK 0xFF00
+#define KCI_RETURN_CORE_LIST_SHIFT 8
+#define KCI_RETURN_ERROR_CODE_MASK (BIT(KCI_RETURN_CORE_LIST_SHIFT) - 1u)
+#define KCI_RETURN_GET_CORE_LIST(ret)                                          \
+	((KCI_RETURN_CORE_LIST_MASK & (ret)) >> KCI_RETURN_CORE_LIST_SHIFT)
+#define KCI_RETURN_GET_ERROR_CODE(ret) (KCI_RETURN_ERROR_CODE_MASK & (ret))
+
 #if IS_ENABLED(CONFIG_GXP_TEST)
 char *gxp_work_mode_name = "direct";
 #else
@@ -92,27 +99,29 @@ static int allocate_vmbox(struct gxp_dev *gxp, struct gxp_virtual_device *vd)
 static void release_vmbox(struct gxp_dev *gxp, struct gxp_virtual_device *vd)
 {
 	struct gxp_kci *kci = &(gxp_mcu_of(gxp)->kci);
+	uint core_list;
 	int ret;
 
 	if (vd->client_id < 0)
 		return;
 
 	ret = gxp_kci_release_vmbox(kci, vd->client_id);
-	if (ret) {
-		/*
-		 * TODO(241057541): Remove this conditional branch after the firmware side
-		 * implements handling allocate_vmbox command.
-		 */
-		if (ret == GCIP_KCI_ERROR_UNIMPLEMENTED)
-			dev_info(
-				gxp->dev,
-				"Releasing VMBox is not implemented from the firmware side");
-		else
-			dev_err(gxp->dev,
-				"Failed to release VMBox for client %d: %d",
-				vd->client_id, ret);
+	if (!ret)
+		goto out;
+	if (ret > 0 &&
+	    KCI_RETURN_GET_ERROR_CODE(ret) == GCIP_KCI_ERROR_ABORTED) {
+		core_list = KCI_RETURN_GET_CORE_LIST(ret);
+		dev_err(gxp->dev,
+			"Firmware failed to gracefully release a VMBox for client %d, core_list=%d",
+			vd->client_id, core_list);
+		gxp_vd_invalidate(gxp, vd);
+		gxp_vd_generate_debug_dump(gxp, vd, core_list);
+	} else {
+		dev_err(gxp->dev,
+			"Failed to request releasing VMBox for client %d: %d",
+			vd->client_id, ret);
 	}
-
+out:
 	vd->client_id = -1;
 }
 
@@ -208,7 +217,7 @@ gxp_mcu_platform_before_vd_block_unready(struct gxp_dev *gxp,
 {
 	if (gxp_is_direct_mode(gxp))
 		return;
-	if (vd->client_id < 0 || vd->state == GXP_VD_UNAVAILABLE)
+	if (vd->client_id < 0 || vd->mcu_crashed)
 		return;
 	if (vd->tpu_client_id >= 0)
 		gxp_mcu_unlink_offload_vmbox(gxp, vd, vd->tpu_client_id,
@@ -216,7 +225,7 @@ gxp_mcu_platform_before_vd_block_unready(struct gxp_dev *gxp,
 	release_vmbox(gxp, vd);
 }
 
-static int gxp_mcu_wakelock_after_blk_on(struct gxp_dev *gxp)
+static int gxp_mcu_pm_after_blk_on(struct gxp_dev *gxp)
 {
 	struct gxp_mcu_firmware *mcu_fw = gxp_mcu_firmware_of(gxp);
 
@@ -225,7 +234,7 @@ static int gxp_mcu_wakelock_after_blk_on(struct gxp_dev *gxp)
 	return gxp_mcu_firmware_run(mcu_fw);
 }
 
-static void gxp_mcu_wakelock_before_blk_off(struct gxp_dev *gxp)
+static void gxp_mcu_pm_before_blk_off(struct gxp_dev *gxp)
 {
 	struct gxp_mcu_firmware *mcu_fw = gxp_mcu_firmware_of(gxp);
 
@@ -350,8 +359,8 @@ void gxp_mcu_dev_init(struct gxp_mcu_dev *mcu_dev)
 	gxp->after_vd_block_ready = gxp_mcu_platform_after_vd_block_ready;
 	gxp->before_vd_block_unready = gxp_mcu_platform_before_vd_block_unready;
 	gxp->request_power_states = gxp_mcu_request_power_states;
-	gxp->wakelock_after_blk_on = gxp_mcu_wakelock_after_blk_on;
-	gxp->wakelock_before_blk_off = gxp_mcu_wakelock_before_blk_off;
+	gxp->pm_after_blk_on = gxp_mcu_pm_after_blk_on;
+	gxp->pm_before_blk_off = gxp_mcu_pm_before_blk_off;
 #if HAS_TPU_EXT
 	gxp->after_map_tpu_mbx_queue = gxp_mcu_after_map_tpu_mbx_queue;
 	gxp->before_unmap_tpu_mbx_queue = gxp_mcu_before_unmap_tpu_mbx_queue;
diff --git a/gxp-pm.c b/gxp-pm.c
index 4c30284..146e788 100644
--- a/gxp-pm.c
+++ b/gxp-pm.c
@@ -13,6 +13,8 @@
 #include <linux/workqueue.h>
 #include <soc/google/exynos_pm_qos.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "gxp-bpm.h"
 #include "gxp-client.h"
 #include "gxp-config.h"
@@ -289,6 +291,28 @@ bool gxp_pm_is_blk_down(struct gxp_dev *gxp, uint timeout_ms)
 	return false;
 }
 
+int gxp_pm_blk_reboot(struct gxp_dev *gxp, uint timeout_ms)
+{
+	int ret;
+
+	ret = gxp_pm_blk_off(gxp);
+	if (ret) {
+		dev_err(gxp->dev, "Failed to turn off BLK_AUR (ret=%d)\n", ret);
+		return ret;
+	}
+
+	if (!gxp_pm_is_blk_down(gxp, timeout_ms)) {
+		dev_err(gxp->dev, "BLK_AUR hasn't been turned off");
+		return -EBUSY;
+	}
+
+	ret = gxp_pm_blk_on(gxp);
+	if (ret)
+		dev_err(gxp->dev, "Failed to turn on BLK_AUR (ret=%d)\n", ret);
+
+	return ret;
+}
+
 int gxp_pm_get_blk_switch_count(struct gxp_dev *gxp)
 {
 	int ret;
@@ -703,11 +727,47 @@ int gxp_pm_update_pm_qos(struct gxp_dev *gxp, s32 int_val, s32 mif_val)
 	return gxp_pm_req_pm_qos(gxp, int_val, mif_val);
 }
 
+static int gxp_pm_power_up(void *data)
+{
+	struct gxp_dev *gxp = data;
+	int ret = gxp_pm_blk_on(gxp);
+
+	if (ret) {
+		dev_err(gxp->dev, "Failed to power on BLK_AUR (ret=%d)\n", ret);
+		return ret;
+	}
+
+	if (gxp->pm_after_blk_on) {
+		ret = gxp->pm_after_blk_on(gxp);
+		if (ret) {
+			gxp_pm_blk_off(gxp);
+			return ret;
+		}
+	}
+
+	return 0;
+}
+
+static int gxp_pm_power_down(void *data)
+{
+	struct gxp_dev *gxp = data;
+
+	if (gxp->pm_before_blk_off)
+		gxp->pm_before_blk_off(gxp);
+	return gxp_pm_blk_off(gxp);
+}
+
 int gxp_pm_init(struct gxp_dev *gxp)
 {
 	struct gxp_power_manager *mgr;
 	struct platform_device *pdev =
 		container_of(gxp->dev, struct platform_device, dev);
+	const struct gcip_pm_args args = {
+		.dev = gxp->dev,
+		.data = gxp,
+		.power_up = gxp_pm_power_up,
+		.power_down = gxp_pm_power_down,
+	};
 	struct resource *r;
 	uint i;
 
@@ -715,6 +775,13 @@ int gxp_pm_init(struct gxp_dev *gxp)
 	if (!mgr)
 		return -ENOMEM;
 	mgr->gxp = gxp;
+
+	mgr->pm = gcip_pm_create(&args);
+	if (IS_ERR(mgr->pm)) {
+		devm_kfree(gxp->dev, mgr);
+		return PTR_ERR(mgr->pm);
+	}
+
 	mutex_init(&mgr->pm_lock);
 	mgr->curr_state = AUR_OFF;
 	mgr->curr_memory_state = AUR_MEM_UNDEFINED;
@@ -767,6 +834,8 @@ int gxp_pm_destroy(struct gxp_dev *gxp)
 	if (IS_GXP_TEST && !mgr)
 		return 0;
 
+	gcip_pm_destroy(mgr->pm);
+
 	exynos_pm_qos_remove_request(&mgr->mif_min);
 	exynos_pm_qos_remove_request(&mgr->int_min);
 	pm_runtime_disable(gxp->dev);
diff --git a/gxp-pm.h b/gxp-pm.h
index 334435e..217c3df 100644
--- a/gxp-pm.h
+++ b/gxp-pm.h
@@ -9,6 +9,8 @@
 
 #include <soc/google/exynos_pm_qos.h>
 
+#include <gcip/gcip-pm.h>
+
 #include "gxp-internal.h"
 
 #define AUR_DVFS_MIN_RATE AUR_UUD_RATE
@@ -105,6 +107,7 @@ static const struct gxp_power_states uud_states = { AUR_UUD, AUR_MEM_UNDEFINED,
 
 struct gxp_power_manager {
 	struct gxp_dev *gxp;
+	struct gcip_pm *pm;
 	struct mutex pm_lock;
 	uint pwr_state_req_count[AUR_NUM_POWER_STATE];
 	uint low_clkmux_pwr_state_req_count[AUR_NUM_POWER_STATE];
@@ -175,6 +178,16 @@ int gxp_pm_blk_off(struct gxp_dev *gxp);
 bool gxp_pm_is_blk_down(struct gxp_dev *gxp, uint timeout_ms);
 
 /**
+ * gxp_pm_blk_reboot() - Reboot the blk.
+ * @gxp: The GXP device to reboot
+ * @timeout_ms: Wait for the block to be turned off for this duration.
+ *
+ * Return:
+ * * 0       - BLK rebooted successfully
+ */
+int gxp_pm_blk_reboot(struct gxp_dev *gxp, uint timeout_ms);
+
+/**
  * gxp_pm_get_blk_state() - Get the blk power state
  * @gxp: The GXP device to sample state
  *
diff --git a/gxp-thermal.c b/gxp-thermal.c
index 812f466..8a9c24b 100644
--- a/gxp-thermal.c
+++ b/gxp-thermal.c
@@ -18,6 +18,8 @@
 #include <linux/thermal.h>
 #include <linux/version.h>
 
+#include <gcip/gcip-pm.h>
+
 /*
  * thermal_cdev_update is moved to drivers/thermal/thermal_core.h in kernel
  * 5.12. The symbol is still exported, manually declare the function prototype
@@ -35,7 +37,6 @@ void thermal_cdev_update(struct thermal_cooling_device *cdev);
 #if GXP_HAS_MCU
 #include "gxp-kci.h"
 #include "gxp-mcu.h"
-#include "gxp-wakelock.h"
 #endif /* GXP_HAS_MCU */
 
 /*
@@ -96,7 +97,8 @@ static int gxp_set_cur_state(struct thermal_cooling_device *cdev,
 #if GXP_HAS_MCU
 			struct gxp_mcu *mcu = gxp_mcu_of(gxp);
 
-			ret = gxp_wakelock_acquire_if_powered(mcu->gxp);
+			ret = gcip_pm_get_if_powered(mcu->gxp->power_mgr->pm,
+						     false);
 			if (ret) {
 				dev_err(dev,
 					"Can't acquire wakelock when powered down: %d\n",
@@ -105,7 +107,7 @@ static int gxp_set_cur_state(struct thermal_cooling_device *cdev,
 			}
 
 			ret = gxp_kci_notify_throttling(&mcu->kci, pwr_state);
-			gxp_wakelock_release(gxp);
+			gcip_pm_put(gxp->power_mgr->pm);
 #endif /* GXP_HAS_MCU */
 		} else {
 			ret = gxp_pm_blk_set_rate_acpm(
diff --git a/gxp-vd.c b/gxp-vd.c
index b4dcad4..fa4ea9e 100644
--- a/gxp-vd.c
+++ b/gxp-vd.c
@@ -22,6 +22,7 @@
 #include "gxp-dma.h"
 #include "gxp-domain-pool.h"
 #include "gxp-doorbell.h"
+#include "gxp-eventfd.h"
 #include "gxp-firmware.h"
 #include "gxp-firmware-data.h"
 #include "gxp-host-device-structs.h"
@@ -31,7 +32,6 @@
 #include "gxp-notification.h"
 #include "gxp-pm.h"
 #include "gxp-vd.h"
-#include "gxp-wakelock.h"
 
 static inline void hold_core_in_reset(struct gxp_dev *gxp, uint core)
 {
@@ -147,8 +147,8 @@ static int map_resource(struct gxp_virtual_device *vd,
 {
 	if (res->daddr == 0)
 		return 0;
-	return gxp_iommu_map(vd->gxp, vd->domain, res->daddr, res->paddr, res->size,
-			     IOMMU_READ | IOMMU_WRITE);
+	return gxp_iommu_map(vd->gxp, vd->domain, res->daddr, res->paddr,
+			     res->size, IOMMU_READ | IOMMU_WRITE);
 }
 
 /* Reverts map_resource. */
@@ -382,7 +382,8 @@ static int alloc_and_map_fw_image(struct gxp_dev *gxp,
 	/* Maps all FW regions together and no rwdata_sgt in this case. */
 	if (ro_size == gxp->fwbufs[0].size)
 		return gxp_iommu_map(gxp, gdomain, gxp->fwbufs[0].daddr,
-				     gxp->fwbufs[0].paddr, ro_size * GXP_NUM_CORES,
+				     gxp->fwbufs[0].paddr,
+				     ro_size * GXP_NUM_CORES,
 				     IOMMU_READ | IOMMU_WRITE);
 
 	dev_info(gxp->dev, "mapping firmware RO size %#zx", ro_size);
@@ -413,7 +414,8 @@ static int alloc_and_map_fw_image(struct gxp_dev *gxp,
 		if (ret) {
 			dev_err(gxp->dev, "map firmware RW for core %d failed",
 				i);
-			gxp_iommu_unmap(gxp, gdomain, gxp->fwbufs[i].daddr, ro_size);
+			gxp_iommu_unmap(gxp, gdomain, gxp->fwbufs[i].daddr,
+					ro_size);
 			goto err_unmap;
 		}
 	}
@@ -443,7 +445,8 @@ static void unmap_and_free_fw_image(struct gxp_dev *gxp,
 	int i;
 
 	if (ro_size == gxp->fwbufs[0].size) {
-		gxp_iommu_unmap(gxp, gdomain, gxp->fwbufs[0].daddr, ro_size * GXP_NUM_CORES);
+		gxp_iommu_unmap(gxp, gdomain, gxp->fwbufs[0].daddr,
+				ro_size * GXP_NUM_CORES);
 		return;
 	}
 
@@ -839,6 +842,11 @@ void gxp_vd_release(struct gxp_virtual_device *vd)
 	if (vd->slice_index >= 0)
 		ida_free(&vd->gxp->shared_slice_idp, vd->slice_index);
 	gxp_domain_pool_free(vd->gxp->domain_pool, vd->domain);
+
+	if (vd->invalidate_eventfd)
+		gxp_eventfd_put(vd->invalidate_eventfd);
+	vd->invalidate_eventfd = NULL;
+
 	vd->state = GXP_VD_RELEASED;
 	debug_dump_unlock(vd);
 	gxp_vd_put(vd);
@@ -847,23 +855,24 @@ void gxp_vd_release(struct gxp_virtual_device *vd)
 int gxp_vd_block_ready(struct gxp_virtual_device *vd)
 {
 	struct gxp_dev *gxp = vd->gxp;
+	enum gxp_virtual_device_state orig_state;
 	int ret;
 
 	lockdep_assert_held_write(&gxp->vd_semaphore);
 
-	if (vd->state == GXP_VD_SUSPENDED)
-		return 0;
-	if (vd->state != GXP_VD_OFF)
+	orig_state = vd->state;
+	if (orig_state != GXP_VD_OFF && orig_state != GXP_VD_SUSPENDED)
 		return -EINVAL;
 	ret = gxp_dma_domain_attach_device(gxp, vd->domain, vd->core_list);
 	if (ret)
 		return ret;
-	vd->state = GXP_VD_READY;
+	if (orig_state == GXP_VD_OFF)
+		vd->state = GXP_VD_READY;
 	if (gxp->after_vd_block_ready) {
 		ret = gxp->after_vd_block_ready(gxp, vd);
 		if (ret) {
 			gxp_dma_domain_detach_device(gxp, vd->domain);
-			vd->state = GXP_VD_OFF;
+			vd->state = orig_state;
 			return ret;
 		}
 	}
@@ -874,8 +883,11 @@ void gxp_vd_block_unready(struct gxp_virtual_device *vd)
 {
 	struct gxp_dev *gxp = vd->gxp;
 
+	lockdep_assert_held_write(&gxp->vd_semaphore);
+
 	if (gxp->before_vd_block_unready)
 		gxp->before_vd_block_unready(gxp, vd);
+	gxp_dma_domain_detach_device(gxp, vd->domain);
 }
 
 int gxp_vd_run(struct gxp_virtual_device *vd)
@@ -974,10 +986,8 @@ void gxp_vd_stop(struct gxp_virtual_device *vd)
 	}
 
 	gxp_firmware_stop(gxp, vd, core_list);
-	if (vd->state == GXP_VD_READY || vd->state == GXP_VD_RUNNING ||
-	    vd->state == GXP_VD_UNAVAILABLE)
-		gxp_dma_domain_detach_device(gxp, vd->domain);
-	vd->state = GXP_VD_OFF;
+	if (vd->state != GXP_VD_UNAVAILABLE)
+		vd->state = GXP_VD_OFF;
 
 	debug_dump_unlock(vd);
 }
@@ -1030,19 +1040,22 @@ void gxp_vd_suspend(struct gxp_virtual_device *vd)
 	u32 boot_state;
 	uint failed_cores = 0;
 
-#ifdef DISABLE_VD_SUSPEND_RESUME_SUPPORT
-	if (!gxp_is_direct_mode(gxp))
+	if (!gxp_is_direct_mode(gxp) && gxp_core_boot)
 		return gxp_vd_stop(vd);
-#endif
 	lockdep_assert_held_write(&gxp->vd_semaphore);
 	debug_dump_lock(gxp, vd);
 
-	dev_info(gxp->dev, "Suspending VD ...\n");
+	dev_info(gxp->dev, "Suspending VD vdid=%d client_id=%d...\n", vd->vdid,
+		 vd->client_id);
 	if (vd->state == GXP_VD_SUSPENDED) {
 		dev_err(gxp->dev,
 			"Attempt to suspend a virtual device twice\n");
 		goto out;
 	}
+	if (!gxp_core_boot) {
+		vd->state = GXP_VD_SUSPENDED;
+		goto out;
+	}
 	gxp_pm_force_clkmux_normal(gxp);
 	/*
 	 * Start the suspend process for all of this VD's cores without waiting
@@ -1108,7 +1121,6 @@ void gxp_vd_suspend(struct gxp_virtual_device *vd)
 					     LPM_CG_STATE);
 		}
 	}
-	gxp_dma_domain_detach_device(gxp, vd->domain);
 	if (vd->state == GXP_VD_UNAVAILABLE) {
 		/* shutdown all cores if virtual device is unavailable */
 		for (phys_core = 0; phys_core < GXP_NUM_CORES; phys_core++)
@@ -1142,20 +1154,24 @@ int gxp_vd_resume(struct gxp_virtual_device *vd)
 
 	lockdep_assert_held_write(&gxp->vd_semaphore);
 	debug_dump_lock(gxp, vd);
-	dev_info(gxp->dev, "Resuming VD ...\n");
+	dev_info(gxp->dev, "Resuming VD vdid=%d client_id=%d...\n", vd->vdid,
+		 vd->client_id);
 	if (vd->state != GXP_VD_SUSPENDED) {
 		dev_err(gxp->dev,
 			"Attempt to resume a virtual device which was not suspended\n");
 		ret = -EBUSY;
 		goto out;
 	}
+	if (!gxp_core_boot) {
+		vd->state = GXP_VD_RUNNING;
+		goto out;
+	}
 	gxp_pm_force_clkmux_normal(gxp);
 	curr_blk_switch_count = gxp_pm_get_blk_switch_count(gxp);
 
 	/* Restore the doorbells state for this VD. */
 	vd_restore_doorbells(vd);
 
-	gxp_dma_domain_attach_device(gxp, vd->domain, core_list);
 	/*
 	 * Start the resume process for all of this VD's cores without waiting
 	 * for completion.
@@ -1237,7 +1253,6 @@ int gxp_vd_resume(struct gxp_virtual_device *vd)
 			if (core_list & BIT(phys_core))
 				gxp_pm_core_off(gxp, phys_core);
 		}
-		gxp_dma_domain_detach_device(gxp, vd->domain);
 	} else {
 		vd->state = GXP_VD_RUNNING;
 	}
@@ -1442,13 +1457,10 @@ void gxp_vd_put(struct gxp_virtual_device *vd)
 		kfree(vd);
 }
 
-void gxp_vd_invalidate(struct gxp_dev *gxp, int client_id, uint core_list)
+void gxp_vd_invalidate_with_client_id(struct gxp_dev *gxp, int client_id,
+				      uint core_list)
 {
 	struct gxp_client *client = NULL, *c;
-	struct gxp_virtual_device *vd;
-	release_unconsumed_async_resps_t release_unconsumed_async_resps =
-		gxp->mailbox_mgr->release_unconsumed_async_resps;
-	int ret;
 	uint core;
 
 	/*
@@ -1494,51 +1506,69 @@ void gxp_vd_invalidate(struct gxp_dev *gxp, int client_id, uint core_list)
 		return;
 	}
 
-	dev_err(gxp->dev, "Invalidate a VD, VDID=%d, client_id=%d",
-		client->vd->vdid, client_id);
+	gxp_vd_invalidate(gxp, client->vd);
+	gxp_vd_generate_debug_dump(gxp, client->vd, core_list);
 
-	if (client->vd->state != GXP_VD_UNAVAILABLE) {
-		if (client->has_block_wakelock) {
-			if (release_unconsumed_async_resps)
-				release_unconsumed_async_resps(client->vd);
-			gxp_vd_block_unready(client->vd);
-		}
+	up_write(&gxp->vd_semaphore);
+	up_write(&client->semaphore);
+}
+
+void gxp_vd_invalidate(struct gxp_dev *gxp, struct gxp_virtual_device *vd)
+{
+	lockdep_assert_held_write(&gxp->vd_semaphore);
+
+	dev_err(gxp->dev, "Invalidate a VD, VDID=%d, client_id=%d", vd->vdid,
+		vd->client_id);
 
-		client->vd->state = GXP_VD_UNAVAILABLE;
-		if (client->vd_invalid_eventfd)
-			gxp_eventfd_signal(client->vd_invalid_eventfd);
+	if (vd->state != GXP_VD_UNAVAILABLE) {
+		if (gxp->mailbox_mgr->release_unconsumed_async_resps)
+			gxp->mailbox_mgr->release_unconsumed_async_resps(vd);
+
+		vd->state = GXP_VD_UNAVAILABLE;
+
+		if (vd->invalidate_eventfd)
+			gxp_eventfd_signal(vd->invalidate_eventfd);
 	} else {
 		dev_dbg(gxp->dev, "This VD is already invalidated");
 	}
+}
+
+void gxp_vd_generate_debug_dump(struct gxp_dev *gxp,
+				struct gxp_virtual_device *vd, uint core_list)
+{
+	int ret;
+
+	if (!gxp_debug_dump_is_enabled() || !core_list)
+		return;
+
+	lockdep_assert_held_write(&gxp->vd_semaphore);
 
 	/*
 	 * We should increase the refcount of @vd because @gxp->vd_semaphore will be
 	 * released below and the client can release it asynchronously.
 	 */
-	vd = gxp_vd_get(client->vd);
+	vd = gxp_vd_get(vd);
 
 	/*
-	 * Release @gxp->vd_semaphore before generating a debug dump not to block other
-	 * virtual devices proceeding their work.
+	 * Release @gxp->vd_semaphore before generating a debug dump and hold it
+	 * again after completing debug dump to not block other virtual devices
+	 * proceeding their work.
 	 */
 	up_write(&gxp->vd_semaphore);
-	up_write(&client->semaphore);
 	mutex_lock(&vd->debug_dump_lock);
 
 	/*
 	 * Process debug dump if its enabled and core_list is not empty.
 	 * Keep on hold the client lock while processing the dumps. vd
-	 * lock would be taken  and released inside the debug dump
+	 * lock would be taken and released inside the debug dump
 	 * implementation logic ahead.
 	 */
-	if (gxp_debug_dump_is_enabled() && core_list != 0) {
-		ret = gxp_debug_dump_process_dump_mcu_mode(gxp, core_list, vd);
-		if (ret)
-			dev_err(gxp->dev,
-				"debug dump processing failed (ret=%d).\n",
-				ret);
-	}
+	ret = gxp_debug_dump_process_dump_mcu_mode(gxp, core_list, vd);
+	if (ret)
+		dev_err(gxp->dev, "debug dump processing failed (ret=%d).\n",
+			ret);
 
 	mutex_unlock(&vd->debug_dump_lock);
+	down_write(&gxp->vd_semaphore);
 	gxp_vd_put(vd);
 }
diff --git a/gxp-vd.h b/gxp-vd.h
index c7fd4e4..e0bd5ce 100644
--- a/gxp-vd.h
+++ b/gxp-vd.h
@@ -151,6 +151,13 @@ struct gxp_virtual_device {
 	struct list_head gxp_fence_list;
 	/* Protects changing the state of vd while generating a debug dump. */
 	struct mutex debug_dump_lock;
+	/* An eventfd which will be triggered when this vd is invalidated. */
+	struct gxp_eventfd *invalidate_eventfd;
+	/*
+	 * If true, the MCU FW communicating with this VD has been crashed and it must not work
+	 * with any MCU FW anymore regardless of its state.
+	 */
+	bool mcu_crashed;
 };
 
 /*
@@ -335,6 +342,8 @@ int gxp_vd_resume(struct gxp_virtual_device *vd);
  * The state of @vd should be GXP_VD_OFF before calling this function.
  * If this function runs successfully, the state becomes GXP_VD_READY.
  *
+ * The caller must have locked gxp->vd_semaphore for writing.
+ *
  * Return:
  * * 0          - Success
  * * -EINVAL    - The VD is not in GXP_VD_OFF state
@@ -343,14 +352,16 @@ int gxp_vd_resume(struct gxp_virtual_device *vd);
 int gxp_vd_block_ready(struct gxp_virtual_device *vd);
 
 /**
- * gxp_vd_block_unready() - This is called before one or both of the virtual device and block
- * wakelock is going to be released.
+ * gxp_vd_block_unready() - This is called before the block wakelock is going to be released.
  *
  * @vd: The virtual device to release the resources
  *
  * This function must be called only when the client holds the block wakelock and allocated a
  * virtual device. It doesn't have a dependency on the state of @vd, but also doesn't change the
- * state.
+ * state in normal situation. However, if an unexpected error happens, the state can be changed
+ * to GXP_VD_UNAVAILABLE.
+ *
+ * The caller must have locked gxp->vd_semaphore for writing.
  */
 void gxp_vd_block_unready(struct gxp_virtual_device *vd);
 
@@ -383,18 +394,47 @@ void gxp_vd_put(struct gxp_virtual_device *vd);
 
 /*
  * Change the status of the vd of @client_id to GXP_VD_UNAVAILABLE.
- * Internally, it will discard all pending/unconsumed user commands
- * and call the `gxp_vd_block_unready` function.
+ * Internally, it will discard all pending/unconsumed user commands and call the
+ * `gxp_vd_block_unready` function.
  *
- * This function will be called when the `CLIENT_FATAL_ERROR_NOTIFY`
- * RKCI has been sent from the firmware side.
+ * This function will be called when the `CLIENT_FATAL_ERROR_NOTIFY` RKCI has been sent from the
+ * firmware side.
  *
  * @gxp: The GXP device to obtain the handler for
  * @client_id: client_id of the crashed vd.
- * @core_list: A bitfield enumerating the physical cores on which
- *             crash is reported from firmware.
+ * @core_list: A bitfield enumerating the physical cores on which crash is reported from firmware.
+ */
+void gxp_vd_invalidate_with_client_id(struct gxp_dev *gxp, int client_id,
+				      uint core_list);
+
+/*
+ * Changes the status of the @vd to GXP_VD_UNAVAILABLE.
+ * Internally, it will discard all pending/unconsumed user commands.
+ *
+ * This function will be called when some unexpected errors happened and cannot proceed requests
+ * anymore with this @vd.
+ *
+ * The caller must have locked gxp->vd_semaphore for writing.
+ *
+ * @gxp: The GXP device to obtain the handler for.
+ * @vd: The virtual device to be invaliated.
+ */
+void gxp_vd_invalidate(struct gxp_dev *gxp, struct gxp_virtual_device *vd);
+
+/*
+ * Generates a debug dump of @vd which utilizes @core_list cores.
+ *
+ * This function is usually called in the MCU mode that the kernel driver cannot decide which cores
+ * will be used by @vd.
+ *
+ * The caller must have locked gxp->vd_semaphore for writing.
+ *
+ * @gxp: The GXP device to obtain the handler for.
+ * @vd: The virtual device to be dumped.
+ * @core_list: A bitfield enumerating the physical cores on which crash is reported from firmware.
  */
-void gxp_vd_invalidate(struct gxp_dev *gxp, int client_id, uint core_list);
+void gxp_vd_generate_debug_dump(struct gxp_dev *gxp,
+				struct gxp_virtual_device *vd, uint core_list);
 
 /*
  * An ID between 0~GXP_NUM_CORES-1 and is unique to each VD.
diff --git a/gxp-wakelock.c b/gxp-wakelock.c
deleted file mode 100644
index ba4d8d8..0000000
--- a/gxp-wakelock.c
+++ /dev/null
@@ -1,183 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0
-/*
- * GXP wakelock support
- *
- * Copyright (C) 2022 Google LLC
- */
-
-#include "gxp-client.h"
-#include "gxp-dma.h"
-#include "gxp-pm.h"
-#include "gxp-wakelock.h"
-
-int gxp_wakelock_init(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr;
-
-	mgr = devm_kzalloc(gxp->dev, sizeof(*mgr), GFP_KERNEL);
-	if (!mgr)
-		return -ENOMEM;
-
-	mutex_init(&mgr->lock);
-
-	gxp->wakelock_mgr = mgr;
-
-	return 0;
-}
-
-static int gxp_wakelock_acquire_locked(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-	int ret = 0;
-
-	if (mgr->suspended) {
-		/*
-		 * Don't allow a new client to obtain a wakelock, powering up
-		 * BLK_AUR, when the system is going to sleep.
-		 */
-		dev_warn(gxp->dev,
-			 "Attempt to obtain wakelock while suspending.\n");
-		ret = -EAGAIN;
-		goto out;
-	}
-
-	if (!mgr->count++) {
-		ret = gxp_pm_blk_on(gxp);
-		if (ret) {
-			dev_err(gxp->dev,
-				"Failed to power on BLK_AUR (ret=%d, client count=%u)\n",
-				ret, mgr->count);
-			goto err_blk_on;
-		}
-		if (gxp->wakelock_after_blk_on) {
-			ret = gxp->wakelock_after_blk_on(gxp);
-			if (ret) {
-				gxp_pm_blk_off(gxp);
-				goto err_blk_on;
-			}
-		}
-	}
-
-out:
-	return ret;
-
-err_blk_on:
-	mgr->count--;
-	return ret;
-}
-
-int gxp_wakelock_acquire(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-	int ret;
-
-	mutex_lock(&mgr->lock);
-	ret = gxp_wakelock_acquire_locked(gxp);
-	mutex_unlock(&mgr->lock);
-
-	return ret;
-}
-
-int gxp_wakelock_acquire_if_powered(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-	int ret = -EAGAIN;
-
-	mutex_lock(&mgr->lock);
-	if (mgr->count)
-		ret = gxp_wakelock_acquire_locked(gxp);
-	mutex_unlock(&mgr->lock);
-
-	return ret;
-}
-
-void gxp_wakelock_release(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-	int ret = 0;
-
-	mutex_lock(&mgr->lock);
-
-	if (!mgr->count) {
-		dev_err(gxp->dev,
-			"Attempt to release wakelock with none held.\n");
-		goto out;
-	}
-
-	if (!--mgr->count) {
-		if (gxp->wakelock_before_blk_off)
-			gxp->wakelock_before_blk_off(gxp);
-		ret = gxp_pm_blk_off(gxp);
-		if (ret)
-			dev_err(gxp->dev,
-				"Failed to power down BLK_AUR (ret=%d, client count=%u)\n",
-				ret, mgr->count);
-	}
-
-out:
-	mutex_unlock(&mgr->lock);
-}
-
-int gxp_wakelock_suspend(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-	int ret;
-	struct gxp_client *client;
-
-	if (!mutex_trylock(&mgr->lock))
-		return -EAGAIN;
-
-	/* Can't suspend if there are any active clients */
-	mgr->suspended = mgr->count == 0;
-	ret = mgr->suspended ? 0 : -EAGAIN;
-
-	/* Suspend successful. Can exit now. */
-	if (!ret)
-		goto out;
-
-	/* Log clients currently holding a wakelock */
-	if (!mutex_trylock(&gxp->client_list_lock)) {
-		dev_warn_ratelimited(
-			gxp->dev,
-			"Unable to get client list lock on suspend failure\n");
-		goto out;
-	}
-
-	list_for_each_entry(client, &gxp->client_list, list_entry) {
-		if (!down_read_trylock(&client->semaphore)) {
-			dev_warn_ratelimited(
-				gxp->dev,
-				"Unable to acquire client lock (tgid=%d pid=%d)\n",
-				client->tgid, client->pid);
-			continue;
-		}
-
-		if (client->has_block_wakelock)
-			dev_warn_ratelimited(
-				gxp->dev,
-				"Cannot suspend with client holding wakelock (tgid=%d pid=%d)\n",
-				client->tgid, client->pid);
-
-		up_read(&client->semaphore);
-	}
-
-	mutex_unlock(&gxp->client_list_lock);
-
-out:
-	mutex_unlock(&mgr->lock);
-
-	return ret;
-}
-
-int gxp_wakelock_resume(struct gxp_dev *gxp)
-{
-	struct gxp_wakelock_manager *mgr = gxp->wakelock_mgr;
-
-	mutex_lock(&mgr->lock);
-
-	mgr->suspended = false;
-
-	mutex_unlock(&mgr->lock);
-
-	return 0;
-}
diff --git a/gxp-wakelock.h b/gxp-wakelock.h
deleted file mode 100644
index e02bdb4..0000000
--- a/gxp-wakelock.h
+++ /dev/null
@@ -1,88 +0,0 @@
-/* SPDX-License-Identifier: GPL-2.0 */
-/*
- * GXP wakelock support
- *
- * Copyright (C) 2022 Google LLC
- */
-#ifndef __GXP_WAKELOCK_H__
-#define __GXP_WAKELOCK_H__
-
-#include "gxp-internal.h"
-#include "gxp.h"
-
-struct gxp_wakelock_manager {
-	/* Protects count and suspended */
-	struct mutex lock;
-	uint count;
-	bool suspended;
-};
-
-/**
- * gxp_wakelock_init() - Initialize wakelock support
- * @gxp: The GXP device to initialize wakelock support for
- *
- * Return:
- * * 0       - Success
- * * -ENOMEM - Insufficient memory is available to initialize support
- */
-int gxp_wakelock_init(struct gxp_dev *gxp);
-
-/**
- * gxp_wakelock_acquire() - Increment the GXP wakelock counter
- * @gxp: The GXP device to increment the wakelock counter for
- *
- * If the wakelock counter transitions from 0 to 1, this will result in BLK_AUR
- * being powered on.
- *
- * Return:
- * * 0       - Success
- * * -EAGAIN - The system is suspending and BLK_AUR cannot be powered on
- * * Other   - An attempt to power on BLK_AUR failed
- */
-int gxp_wakelock_acquire(struct gxp_dev *gxp);
-
-/**
- * gxp_wakelock_acquire_if_powered() - Increment the GXP wakelock counter if
- *                                     the counter is nonzero.
- * @gxp: The GXP device to increment the wakelock counter for
- *
- * Similar to gxp_wakelock_acquire, but only increment the wakelock counter if
- * the counter is nonzero.
- *
- * Return:
- * * 0       - Success
- * * -EAGAIN - Wakelock counter is zero
- * * Other   - Error returned by gxp_wakelock_acquire
- */
-int gxp_wakelock_acquire_if_powered(struct gxp_dev *gxp);
-
-/**
- * gxp_wakelock_release() - Decrement the GXP wakelock counter
- * @gxp: The GXP device to decrement the wakelock counter for
- *
- * If the wakelock counter transitions from 1 to 0, this will result in BLK_AUR
- * being powered off. In the event BLK_AUR cannot be powered off, a message
- * will be logged, but the wakelock will still be released.
- */
-void gxp_wakelock_release(struct gxp_dev *gxp);
-
-/**
- * gxp_wakelock_suspend() - Check if the wakelock will allow a system suspend
- * @gxp: The GXP device to check the wakelock of
- *
- * Return:
- * * 0       - The wakelock has been suspended and is ready for system suspend
- * * -EAGAIN - The wakelock is held, and system suspend should be aborted
- */
-int gxp_wakelock_suspend(struct gxp_dev *gxp);
-
-/**
- * gxp_wakelock_resume() - Notify the wakelock that system suspend has exited
- * @gxp: The GXP device to notify the wakelock of
- *
- * Return:
- * * 0 - The wakelock is ready to be acquired again
- */
-int gxp_wakelock_resume(struct gxp_dev *gxp);
-
-#endif /* __GXP_WAKELOCK_H__ */
diff --git a/gxp.h b/gxp.h
index 5bb868a..b5c64f5 100644
--- a/gxp.h
+++ b/gxp.h
@@ -13,7 +13,7 @@
 
 /* Interface Version */
 #define GXP_INTERFACE_VERSION_MAJOR 1
-#define GXP_INTERFACE_VERSION_MINOR 9
+#define GXP_INTERFACE_VERSION_MINOR 10
 #define GXP_INTERFACE_VERSION_BUILD 0
 
 /*