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

* whitechapel: (36 commits)
  edgetpu: add accessors for firmware build CL and timestamp
  edgetpu: retrieve build CL and timestamp from firmware via KCI
  edgetpu: pin user pages before holding locks
  edgetpu: remove edgetpu_ioctl_check_group
  edgetpu: dmabuf ioctls accept client as param
  edgetpu: abrolhos ignore return value of OPEN_DEVICE
  edgetpu: multi ioctl handlers accept client param
  edgetpu: ioctl (un)set eventfd accept client as param
  edgetpu: add launch_complete firmware callback
  edgetpu: ioctl_(un)map_buffer accept client as param
  ...

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

25 files changed, 538 insertions, 269 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index b364f41..bf17cf0 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -12,8 +12,8 @@ endif
 
 edgetpu-fw-objs := edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-shared-fw.o
 edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o $(edgetpu-fw-objs)
-abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-direct.o abrolhos-core.o abrolhos-platform.o abrolhos-iommu.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o $(edgetpu-objs)
-CFLAGS_abrolhos-direct.o := -DCONFIG_ABROLHOS=1
+abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-iommu.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o $(edgetpu-objs)
+CFLAGS_abrolhos-fs.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-core.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device-group.o := -DCONFIG_ABROLHOS=1
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index f57050a..58b9835 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -16,7 +16,7 @@ endif
 edgetpu-fw-objs := edgetpu-firmware-util.o edgetpu-shared-fw.o edgetpu-firmware.o
 edgetpu-objs	:= edgetpu-core.o edgetpu-mailbox.o edgetpu-kci.o edgetpu-device-group.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o $(edgetpu-fw-objs)
 
-abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-direct.o abrolhos-platform.o abrolhos-iommu.o abrolhos-thermal.o abrolhos-pm.o $(edgetpu-objs)
+abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-fs.o abrolhos-platform.o abrolhos-iommu.o abrolhos-thermal.o abrolhos-pm.o $(edgetpu-objs)
 
 KBUILD_OPTIONS += CONFIG_ABROLHOS=m
 
diff --git a/drivers/edgetpu/abrolhos-firmware.c b/drivers/edgetpu/abrolhos-firmware.c
index 649beb8..07e6792 100644
--- a/drivers/edgetpu/abrolhos-firmware.c
+++ b/drivers/edgetpu/abrolhos-firmware.c
@@ -7,7 +7,6 @@
 
 #include <linux/dma-mapping.h>
 #include <linux/gsa/gsa_tpu.h>
-#include <linux/sizes.h>
 #include <linux/slab.h>
 
 #include "abrolhos-firmware.h"
@@ -18,8 +17,6 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 
-#define ABROLHOS_FW_HEADER_SIZE		SZ_4K
-
 static int abrolhos_firmware_alloc_buffer(
 		struct edgetpu_firmware *et_fw,
 		struct edgetpu_firmware_buffer *fw_buf)
diff --git a/drivers/edgetpu/abrolhos-firmware.h b/drivers/edgetpu/abrolhos-firmware.h
index 7deb36f..8e7db77 100644
--- a/drivers/edgetpu/abrolhos-firmware.h
+++ b/drivers/edgetpu/abrolhos-firmware.h
@@ -7,9 +7,13 @@
 #ifndef __ABROLHOS_FIRMWARE_H__
 #define __ABROLHOS_FIRMWARE_H__
 
+#include <linux/sizes.h>
+
 #include "edgetpu-internal.h"
 #include "edgetpu.h"
 
+/* abrolhos FW header size */
+#define ABROLHOS_FW_HEADER_SIZE SZ_4K
 /* The offset to the signed firmware header. */
 #define ABROLHOS_HEADER_OFFSET 0x400
 /* The offset to image configuration. */
@@ -25,6 +29,23 @@ struct abrolhos_image_config {
 	struct edgetpu_fw_version firmware_versions;
 } __packed;
 
+/*
+ * Abrolhos firmware header.
+ */
+struct abrolhos_image_header {
+	char sig[512];
+	char pub[512];
+	int Magic;
+	int Generation;
+	int RollbackInfo;
+	int Length;
+	char Flags[16];
+	char BodyHash[32];
+	char ChipId[32];
+	char AuthConfig[256];
+	struct abrolhos_image_config ImageConfig;
+};
+
 int abrolhos_edgetpu_firmware_create(struct edgetpu_dev *etdev);
 void abrolhos_edgetpu_firmware_destroy(struct edgetpu_dev *etdev);
 
diff --git a/drivers/edgetpu/abrolhos-direct.c b/drivers/edgetpu/abrolhos-fs.c
index 6caef78..5babc06 100644
--- a/drivers/edgetpu/abrolhos-direct.c
+++ b/drivers/edgetpu/abrolhos-fs.c
@@ -1,2 +1,2 @@
 // SPDX-License-Identifier: GPL-2.0
-#include "edgetpu-direct.c"
+#include "edgetpu-fs.c"
diff --git a/drivers/edgetpu/abrolhos-platform.c b/drivers/edgetpu/abrolhos-platform.c
index 8f40c8e..b896865 100644
--- a/drivers/edgetpu/abrolhos-platform.c
+++ b/drivers/edgetpu/abrolhos-platform.c
@@ -403,8 +403,8 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 
 out:
 	dev_dbg(dev, "Probe finished, powering down\n");
-	/* Turn the device off until a client request is received */
-	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev);
+	/* Turn the device off unless a client request is already received. */
+	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, false);
 
 	return ret;
 out_fw_destroy:
@@ -414,8 +414,7 @@ out_tel_exit:
 out_cleanup_fw:
 	edgetpu_platform_cleanup_fw_region(edgetpu_pdev);
 	dev_dbg(dev, "Probe finished with error %d, powering down\n", ret);
-	/* Turn the device off until a client request is received */
-	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev);
+	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, true);
 	return ret;
 }
 
@@ -435,7 +434,7 @@ static int edgetpu_platform_remove(struct platform_device *pdev)
 	edgetpu_platform_cleanup_fw_region(edgetpu_pdev);
 	edgetpu_device_remove(etdev);
 	edgetpu_pm_put(etdev->pm);
-	edgetpu_pm_shutdown(etdev);
+	edgetpu_pm_shutdown(etdev, true);
 	abrolhos_pm_destroy(etdev);
 	return 0;
 }
diff --git a/drivers/edgetpu/abrolhos-platform.h b/drivers/edgetpu/abrolhos-platform.h
index 6345da7..d2ee05e 100644
--- a/drivers/edgetpu/abrolhos-platform.h
+++ b/drivers/edgetpu/abrolhos-platform.h
@@ -9,11 +9,15 @@
 
 #include <linux/device.h>
 #include <linux/io.h>
+#include <linux/kernel.h>
 #include <linux/types.h>
 
 #include "edgetpu-internal.h"
 #include "abrolhos-pm.h"
 
+#define to_abrolhos_dev(etdev)                                                 \
+	container_of(etdev, struct edgetpu_platform_dev, edgetpu_dev)
+
 struct edgetpu_platform_pwr {
 	struct mutex policy_lock;
 	enum tpu_pwr_state curr_policy;
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index bfe5242..fed22f2 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -513,7 +513,7 @@ static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
 
 	mutex_init(&edgetpu_pdev->platform_pwr.policy_lock);
 	abrolhos_pwr_debugfs_dir =
-		debugfs_create_dir("power", edgetpu_dev_debugfs_dir());
+		debugfs_create_dir("power", edgetpu_fs_debugfs_dir());
 	debugfs_create_file("state", 0660, abrolhos_pwr_debugfs_dir,
 			dev, &fops_tpu_pwr_state);
 	debugfs_create_file("vdd_tpu", 0660, abrolhos_pwr_debugfs_dir,
diff --git a/drivers/edgetpu/abrolhos-pm.h b/drivers/edgetpu/abrolhos-pm.h
index 5c0a753..864f070 100644
--- a/drivers/edgetpu/abrolhos-pm.h
+++ b/drivers/edgetpu/abrolhos-pm.h
@@ -16,8 +16,10 @@
 
 #else
 
+static unsigned long exynos_acpm_rate;
 static inline int exynos_acpm_set_rate(unsigned int id, unsigned long rate)
 {
+	exynos_acpm_rate = rate;
 	return 0;
 }
 static inline int exynos_acpm_set_init_freq(unsigned int dfs_id,
@@ -28,7 +30,7 @@ static inline int exynos_acpm_set_init_freq(unsigned int dfs_id,
 static inline unsigned long exynos_acpm_get_rate(unsigned int id,
 						 unsigned long dbg_val)
 {
-	return 0;
+	return exynos_acpm_rate;
 }
 static inline int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
 {
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index 2aa0f43..957fe52 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -243,7 +243,7 @@ static int tpu_thermal_init(struct edgetpu_thermal *thermal, struct device *dev)
 
 	thermal->dev = dev;
 	thermal->cooling_root =
-		debugfs_create_dir("cooling", edgetpu_dev_debugfs_dir());
+		debugfs_create_dir("cooling", edgetpu_fs_debugfs_dir());
 
 	err = tpu_thermal_cooling_register(thermal, EDGETPU_COOLING_NAME);
 	if (err) {
diff --git a/drivers/edgetpu/abrolhos/config-tpu-cpu.h b/drivers/edgetpu/abrolhos/config-tpu-cpu.h
index 6060892..81e6187 100644
--- a/drivers/edgetpu/abrolhos/config-tpu-cpu.h
+++ b/drivers/edgetpu/abrolhos/config-tpu-cpu.h
@@ -8,15 +8,17 @@
 #ifndef __ABROLHOS_CONFIG_TPU_CPU_H__
 #define __ABROLHOS_CONFIG_TPU_CPU_H__
 
-/* TODO(137684940): Modify the CSR generator tool to output the values. */
 #define EDGETPU_REG_RESET_CONTROL			0x90010
 #define CPUPORESET					(1 << 1)
+
 #define EDGETPU_REG_INSTRUCTION_REMAP_CONTROL		0x90050
 #define EDGETPU_REG_INSTRUCTION_REMAP_BASE		0x90058
 #define EDGETPU_REG_INSTRUCTION_REMAP_LIMIT		0x90060
 #define EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE		0x90068
 #define EDGETPU_REG_INSTRUCTION_REMAP_SECURITY		0x90070
 #define EDGETPU_REG_SECURITY				0x90048
+
+/* Power Control signals for P-channel interface. */
 #define EDGETPU_REG_POWER_CONTROL			0xA0008
 #define PSTATE						(1 << 0)
 #define PREQ						(1 << 1)
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 8c6fb4d..f1f584b 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -222,9 +222,9 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 	mutex_init(&etdev->state_lock);
 	etdev->state = ETDEV_STATE_NOFW;
 
-	ret = edgetpu_dev_add(etdev);
+	ret = edgetpu_fs_add(etdev);
 	if (ret) {
-		dev_err(etdev->dev, "%s: edgetpu_dev_add returns %d\n",
+		dev_err(etdev->dev, "%s: edgetpu_fs_add returns %d\n",
 			etdev->dev_name, ret);
 		return ret;
 	}
@@ -281,7 +281,7 @@ detach_mmu:
 	edgetpu_mmu_detach(etdev);
 remove_dev:
 	edgetpu_mark_probe_fail(etdev);
-	edgetpu_dev_remove(etdev);
+	edgetpu_fs_remove(etdev);
 	return ret;
 }
 
@@ -290,7 +290,7 @@ void edgetpu_device_remove(struct edgetpu_dev *etdev)
 	edgetpu_chip_exit(etdev);
 	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
 	edgetpu_mmu_detach(etdev);
-	edgetpu_dev_remove(etdev);
+	edgetpu_fs_remove(etdev);
 }
 
 struct edgetpu_client *edgetpu_client_add(struct edgetpu_dev *etdev)
@@ -366,7 +366,7 @@ int __init edgetpu_init(void)
 {
 	int ret;
 
-	ret = edgetpu_dev_init();
+	ret = edgetpu_fs_init();
 	if (ret)
 		return ret;
 	edgetpu_mcp_init();
@@ -376,5 +376,5 @@ int __init edgetpu_init(void)
 void __exit edgetpu_exit(void)
 {
 	edgetpu_mcp_exit();
-	edgetpu_dev_exit();
+	edgetpu_fs_exit();
 }
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index aa5c965..f3e008e 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -95,8 +95,18 @@ static int edgetpu_kci_leave_group_worker(struct kci_worker_param *param)
 static void edgetpu_device_group_kci_leave(struct edgetpu_device_group *group)
 {
 #if IS_ENABLED(CONFIG_ABROLHOS)
+	u8 mailbox_id = group->vii.mailbox->mailbox_id;
+	int ret = edgetpu_kci_close_device(group->etdev->kci, mailbox_id);
+
+	/*
+	 * This should only happen when the FW hasn't driven this KCI, log once
+	 * to prevent log storm.
+	 */
+	if (ret)
+		etdev_warn_once(group->etdev, "Close device failed with %d",
+				ret);
 	return;
-#else
+#else /* !CONFIG_ABROLHOS */
 	struct kci_worker_param *params =
 		kmalloc_array(group->n_clients, sizeof(*params), GFP_KERNEL);
 	struct edgetpu_async_ctx *ctx = edgetpu_async_alloc_ctx();
@@ -125,7 +135,7 @@ static void edgetpu_device_group_kci_leave(struct edgetpu_device_group *group)
 out_free:
 	edgetpu_async_free_ctx(ctx);
 	kfree(params);
-#endif
+#endif /* CONFIG_ABROLHOS */
 }
 
 /*
@@ -137,8 +147,18 @@ static int
 edgetpu_device_group_kci_finalized(struct edgetpu_device_group *group)
 {
 #if IS_ENABLED(CONFIG_ABROLHOS)
+	u8 mailbox_id = group->vii.mailbox->mailbox_id;
+	int ret = edgetpu_kci_open_device(group->etdev->kci, mailbox_id);
+
+	/*
+	 * This should only happen when the FW hasn't driven this KCI, log once
+	 * to prevent log storm.
+	 */
+	if (ret)
+		etdev_warn_once(group->etdev, "Open device failed with %d",
+				ret);
 	return 0;
-#else
+#else /* !CONFIG_ABROLHOS */
 	struct kci_worker_param *params =
 		kmalloc_array(group->n_clients, sizeof(*params), GFP_KERNEL);
 	struct edgetpu_async_ctx *ctx = edgetpu_async_alloc_ctx();
@@ -204,7 +224,7 @@ out_free:
 	edgetpu_async_free_ctx(ctx);
 	kfree(params);
 	return ret;
-#endif
+#endif /* CONFIG_ABROLHOS */
 }
 
 /*
@@ -888,15 +908,18 @@ static void edgetpu_host_map_show(struct edgetpu_mapping *map,
 }
 
 /*
- * Pins the user-space address @host_addr and returns the pinned pages.
+ * Pins the user-space address @arg->host_address and returns the pinned pages.
  * @pnum_pages is set to the number of pages.
  *
  * Returns -errno if failed on pinning @size bytes.
  */
-static struct page **edgetpu_pin_user_pages(
-	struct edgetpu_device_group *group, u64 host_addr, u64 size,
-	enum dma_data_direction dir, uint *pnum_pages)
+static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
+					    struct edgetpu_map_ioctl *arg,
+					    uint *pnum_pages)
 {
+	u64 host_addr = arg->host_address;
+	u64 size = arg->size;
+	const enum dma_data_direction dir = arg->flags & EDGETPU_MAP_DIR_MASK;
 	uint num_pages;
 	ulong offset;
 	struct edgetpu_dev *etdev = group->etdev;
@@ -953,11 +976,10 @@ error:
  */
 static struct edgetpu_host_map *
 alloc_mapping_from_useraddr(struct edgetpu_device_group *group, u64 host_addr,
-			    u64 size, edgetpu_map_flag_t flags)
+			    u64 size, edgetpu_map_flag_t flags,
+			    struct page **pages, uint num_pages)
 {
-	uint num_pages = 0;
 	struct edgetpu_dev *etdev = group->etdev;
-	struct page **pages;
 	struct edgetpu_host_map *hmap;
 	const enum dma_data_direction dir = flags & EDGETPU_MAP_DIR_MASK;
 	int n;
@@ -965,10 +987,6 @@ alloc_mapping_from_useraddr(struct edgetpu_device_group *group, u64 host_addr,
 	int i;
 	int ret;
 
-	pages = edgetpu_pin_user_pages(group, host_addr, size, dir, &num_pages);
-	if (IS_ERR(pages))
-		return (void *)pages;
-
 	hmap = kzalloc(sizeof(*hmap), GFP_KERNEL);
 	if (!hmap) {
 		ret = -ENOMEM;
@@ -1012,7 +1030,6 @@ alloc_mapping_from_useraddr(struct edgetpu_device_group *group, u64 host_addr,
 		}
 	}
 
-	kfree(pages);
 	return hmap;
 
 error_free_sgt:
@@ -1027,7 +1044,6 @@ error_free_sgt:
 error:
 	for (i = 0; i < num_pages; i++)
 		put_page(pages[i]);
-	kfree(pages);
 	if (hmap) {
 		edgetpu_device_group_put(hmap->map.priv);
 		kfree(hmap->sg_tables);
@@ -1105,6 +1121,8 @@ static int group_sync_host_map(struct edgetpu_device_group *group,
 int edgetpu_device_group_map(struct edgetpu_device_group *group,
 			     struct edgetpu_map_ioctl *arg)
 {
+	uint num_pages = 0;
+	struct page **pages;
 	int ret = -EINVAL;
 	u64 host_addr = arg->host_address;
 	u64 size = arg->size;
@@ -1115,6 +1133,11 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 	enum edgetpu_context_id context_id = edgetpu_group_context_id(group);
 	const u32 mmu_flags = map_to_mmu_flags(flags) | EDGETPU_MMU_HOST;
 
+	/* Pin user pages before holding any lock. */
+	pages = edgetpu_pin_user_pages(group, arg, &num_pages);
+	if (IS_ERR(pages))
+		return PTR_ERR(pages);
+
 	mutex_lock(&group->lock);
 	if (!edgetpu_device_group_is_finalized(group)) {
 		ret = -EINVAL;
@@ -1127,7 +1150,8 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 		}
 	}
 
-	hmap = alloc_mapping_from_useraddr(group, host_addr, size, flags);
+	hmap = alloc_mapping_from_useraddr(group, host_addr, size, flags, pages,
+					   num_pages);
 	if (IS_ERR(hmap)) {
 		ret = PTR_ERR(hmap);
 		goto error_unlock_group;
@@ -1175,6 +1199,7 @@ error_release_map:
 
 error_unlock_group:
 	mutex_unlock(&group->lock);
+	kfree(pages);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 7927112..9a652e6 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -52,7 +52,7 @@ struct edgetpu_firmware_private {
 	struct edgetpu_firmware_desc fw_desc;
 	struct edgetpu_firmware_desc bl1_fw_desc;
 	enum edgetpu_firmware_status status;
-	enum edgetpu_fw_flavor fw_flavor;
+	struct edgetpu_fw_info fw_info;
 };
 
 void edgetpu_firmware_set_data(struct edgetpu_firmware *et_fw, void *data)
@@ -252,31 +252,40 @@ static char *fw_flavor_str(enum edgetpu_fw_flavor fw_flavor)
 
 static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 {
+	struct edgetpu_dev *etdev = et_fw->etdev;
 	enum edgetpu_fw_flavor fw_flavor;
 	struct edgetpu_firmware_buffer *fw_buf;
-	struct edgetpu_dev *etdev = et_fw->etdev;
 
 	/* Give the firmware some time to initialize */
 	msleep(100);
-	etdev_dbg(etdev, "Detecting firmware flavor...");
-	fw_flavor = edgetpu_kci_fw_flavor(etdev->kci);
+	etdev_dbg(etdev, "Detecting firmware info...");
+	et_fw->p->fw_info.fw_build_time = 0;
+	et_fw->p->fw_info.fw_flavor = FW_FLAVOR_UNKNOWN;
+	et_fw->p->fw_info.fw_changelist = 0;
+	fw_flavor = edgetpu_kci_fw_info(etdev->kci, &et_fw->p->fw_info);
 	if (fw_flavor < 0) {
 		etdev_err(etdev, "firmware handshake failed: %d", fw_flavor);
 		et_fw->p->status = FW_INVALID;
-		et_fw->p->fw_flavor = FW_FLAVOR_UNKNOWN;
+		et_fw->p->fw_info.fw_flavor = FW_FLAVOR_UNKNOWN;
+		et_fw->p->fw_info.fw_changelist = 0;
+		et_fw->p->fw_info.fw_build_time = 0;
 		return fw_flavor;
 	}
 
 	if (fw_flavor != FW_FLAVOR_BL1) {
 		fw_buf = &et_fw->p->fw_desc.buf;
-		etdev_info(etdev, "loaded %s firmware%s",
+		etdev_info(etdev, "loaded %s firmware%s (%u.%u %u)",
 			   fw_flavor_str(fw_flavor),
-			   fw_buf->flags & FW_ONDEV ? " on device" : "");
+			   fw_buf->flags & FW_ONDEV ? " on device" : "",
+			   etdev->fw_version.major_version,
+			   etdev->fw_version.minor_version,
+			   et_fw->p->fw_info.fw_changelist);
 	} else {
 		etdev_dbg(etdev, "loaded stage 2 bootloader");
 	}
 	et_fw->p->status = FW_VALID;
-	et_fw->p->fw_flavor = fw_flavor;
+	/* In case older firmware that doesn't fill out fw_info. */
+	et_fw->p->fw_info.fw_flavor = fw_flavor;
 	/* Hermosa second-stage bootloader doesn't implement log/trace */
 	if (fw_flavor != FW_FLAVOR_BL1) {
 		int ret = edgetpu_telemetry_kci(etdev);
@@ -290,7 +299,19 @@ static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 enum edgetpu_fw_flavor
 edgetpu_firmware_get_flavor(struct edgetpu_firmware *et_fw)
 {
-	return et_fw->p->fw_flavor;
+	return et_fw->p->fw_info.fw_flavor;
+}
+
+uint32_t
+edgetpu_firmware_get_cl(struct edgetpu_firmware *et_fw)
+{
+	return et_fw->p->fw_info.fw_changelist;
+}
+
+uint64_t
+edgetpu_firmware_get_build_time(struct edgetpu_firmware *et_fw)
+{
+	return et_fw->p->fw_info.fw_build_time;
 }
 
 int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
@@ -335,8 +356,11 @@ int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 	ret = edgetpu_firmware_handshake(et_fw);
 
 	/* Don't start wdt if loaded firmware is second stage bootloader. */
-	if (!ret && !is_bl1_run && et_fw->p->fw_flavor != FW_FLAVOR_BL1)
+	if (!ret && !is_bl1_run && et_fw->p->fw_info.fw_flavor != FW_FLAVOR_BL1)
 		edgetpu_sw_wdt_start(et_fw->etdev);
+
+	if (!ret && !is_bl1_run && handlers && handlers->launch_complete)
+		handlers->launch_complete(et_fw);
 	return ret;
 
 out_unload_new_fw:
@@ -526,14 +550,39 @@ static ssize_t firmware_type_show(
 	if (!et_fw)
 		return -ENODEV;
 	ret = scnprintf(buf, PAGE_SIZE, "%s\n",
-			fw_flavor_str(et_fw->p->fw_flavor));
+			fw_flavor_str(et_fw->p->fw_info.fw_flavor));
 	return ret;
 }
 static DEVICE_ATTR_RO(firmware_type);
 
+static ssize_t firmware_version_show(
+		struct device *dev, struct device_attribute *attr,
+		char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_firmware *et_fw = etdev->firmware;
+	int ret;
+
+	if (!et_fw)
+		return -ENODEV;
+
+	if (etdev->fw_version.kci_version == EDGETPU_INVALID_KCI_VERSION)
+		ret = -ENODATA;
+	else
+		ret = scnprintf(buf, PAGE_SIZE, "%u.%u vii=%u kci=%u cl=%u\n",
+				etdev->fw_version.major_version,
+				etdev->fw_version.minor_version,
+				etdev->fw_version.vii_version,
+				etdev->fw_version.kci_version,
+				et_fw->p->fw_info.fw_changelist);
+	return ret;
+}
+static DEVICE_ATTR_RO(firmware_version);
+
 static struct attribute *dev_attrs[] = {
 	&dev_attr_load_firmware.attr,
 	&dev_attr_firmware_type.attr,
+	&dev_attr_firmware_version.attr,
 	NULL,
 };
 
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index ea9208b..e07ad44 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -29,6 +29,28 @@ enum edgetpu_firmware_status {
 	FW_VALID = 2,
 };
 
+/* Firmware flavors returned via KCI FIRMWARE_INFO command. */
+enum edgetpu_fw_flavor {
+	/* used by host when cannot determine the flavor */
+	FW_FLAVOR_UNKNOWN = 0,
+	/* second-stage bootloader */
+	FW_FLAVOR_BL1 = 1,
+	/* systest app image */
+	FW_FLAVOR_SYSTEST = 2,
+	/* default production app image from DarwiNN team */
+	FW_FLAVOR_PROD_DEFAULT = 3,
+	/* custom image produced by other teams */
+	FW_FLAVOR_CUSTOM = 4,
+};
+
+/* Firmware info filled out via KCI FIRMWARE_INFO command. */
+struct edgetpu_fw_info {
+	uint64_t fw_build_time;		/* BuildData::Timestamp() */
+	uint32_t fw_flavor;		/* enum edgetpu_fw_flavor */
+	uint32_t fw_changelist;		/* BuildData::Changelist() */
+	uint32_t spare[10];
+};
+
 struct edgetpu_firmware_private;
 
 struct edgetpu_firmware {
@@ -121,6 +143,8 @@ struct edgetpu_firmware_handlers {
 	 */
 	int (*prepare_run)(struct edgetpu_firmware *et_fw,
 			   struct edgetpu_firmware_buffer *fw_buf);
+	/* Firmware running, after successful handshake. */
+	void (*launch_complete)(struct edgetpu_firmware *et_fw);
 };
 
 /*
@@ -198,4 +222,10 @@ int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 enum edgetpu_fw_flavor
 edgetpu_firmware_get_flavor(struct edgetpu_firmware *et_fw);
 
+/* Returns the changelist ID of the image loaded on the device. */
+uint32_t edgetpu_firmware_get_cl(struct edgetpu_firmware *et_fw);
+
+/* Returns the build time of the image in seconds since 1970. */
+uint64_t edgetpu_firmware_get_build_time(struct edgetpu_firmware *et_fw);
+
 #endif /* __EDGETPU_FIRMWARE_H__ */
diff --git a/drivers/edgetpu/edgetpu-direct.c b/drivers/edgetpu/edgetpu-fs.c
index e6efe51..c276189 100644
--- a/drivers/edgetpu/edgetpu-direct.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -49,6 +49,22 @@ static atomic_t char_minor = ATOMIC_INIT(-1);
 
 static struct dentry *edgetpu_debugfs_dir;
 
+#define LOCK(client) mutex_lock(&client->group_lock)
+#define UNLOCK(client) mutex_unlock(&client->group_lock)
+/*
+ * Locks @client->group_lock and assigns @client->group to @grp.
+ * Returns -EINVAL if @client is not the leader of the group.
+ */
+#define LOCK_RETURN_IF_NOT_LEADER(client, grp)                                 \
+	do {                                                                   \
+		LOCK(client);                                                  \
+		grp = client->group;                                           \
+		if (!grp || !edgetpu_device_group_is_leader(grp, client)) {    \
+			UNLOCK(client);                                        \
+			return -EINVAL;                                        \
+		}                                                              \
+	} while (0)
+
 int edgetpu_open(struct edgetpu_dev *etdev, struct file *file)
 {
 	struct edgetpu_client *client;
@@ -78,7 +94,7 @@ int edgetpu_open(struct edgetpu_dev *etdev, struct file *file)
 	return 0;
 }
 
-static int edgetpu_dev_open(struct inode *inode, struct file *file)
+static int edgetpu_fs_open(struct inode *inode, struct file *file)
 {
 	struct edgetpu_dev *etdev =
 		container_of(inode->i_cdev, struct edgetpu_dev, cdev);
@@ -86,7 +102,7 @@ static int edgetpu_dev_open(struct inode *inode, struct file *file)
 	return edgetpu_open(etdev, file);
 }
 
-static int etdirect_release(struct inode *inode, struct file *file)
+static int edgetpu_fs_release(struct inode *inode, struct file *file)
 {
 	struct edgetpu_client *client = file->private_data;
 	struct edgetpu_dev *etdev;
@@ -123,21 +139,37 @@ static int etdirect_release(struct inode *inode, struct file *file)
 	return 0;
 }
 
-static int etdirect_set_eventfd(struct edgetpu_device_group *group,
-				struct edgetpu_event_register __user *argp)
+static int edgetpu_ioctl_set_eventfd(struct edgetpu_client *client,
+				     struct edgetpu_event_register __user *argp)
 {
+	struct edgetpu_device_group *group;
+	int ret;
 	struct edgetpu_event_register eventreg;
 
 	if (copy_from_user(&eventreg, argp, sizeof(eventreg)))
 		return -EFAULT;
 
-	return edgetpu_group_set_eventfd(group, eventreg.event_id,
-					 eventreg.eventfd);
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_group_set_eventfd(group, eventreg.event_id,
+					eventreg.eventfd);
+	UNLOCK(client);
+	return ret;
+}
+
+static int edgetpu_ioctl_unset_eventfd(struct edgetpu_client *client,
+				       uint event_id)
+{
+	struct edgetpu_device_group *group;
+
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	edgetpu_group_unset_eventfd(group, event_id);
+	UNLOCK(client);
+	return 0;
 }
 
 static int
-etdirect_set_perdie_eventfd(struct edgetpu_dev *etdev,
-			    struct edgetpu_event_register __user *argp)
+edgetpu_ioctl_set_perdie_eventfd(struct edgetpu_dev *etdev,
+				 struct edgetpu_event_register __user *argp)
 {
 	struct edgetpu_event_register eventreg;
 
@@ -156,8 +188,8 @@ etdirect_set_perdie_eventfd(struct edgetpu_dev *etdev,
 	}
 }
 
-static int etdirect_unset_perdie_eventfd(struct edgetpu_dev *etdev,
-					 uint event_id)
+static int edgetpu_ioctl_unset_perdie_eventfd(struct edgetpu_dev *etdev,
+					      uint event_id)
 {
 	switch (event_id) {
 	case EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE:
@@ -173,7 +205,19 @@ static int etdirect_unset_perdie_eventfd(struct edgetpu_dev *etdev,
 	return 0;
 }
 
-static int etdirect_join_group(struct edgetpu_client *client, u64 leader_fd)
+static int edgetpu_ioctl_finalize_group(struct edgetpu_client *client)
+{
+	struct edgetpu_device_group *group;
+	int ret;
+
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_device_group_finalize(group);
+	UNLOCK(client);
+	return ret;
+}
+
+static int edgetpu_ioctl_join_group(struct edgetpu_client *client,
+				    u64 leader_fd)
 {
 	struct fd f = fdget(leader_fd);
 	struct file *file = f.file;
@@ -224,9 +268,10 @@ static int edgetpu_ioctl_create_group(struct edgetpu_client *client,
 	return 0;
 }
 
-static int etdirect_map_buffer(struct edgetpu_device_group *group,
-			       struct edgetpu_map_ioctl __user *argp)
+static int edgetpu_ioctl_map_buffer(struct edgetpu_client *client,
+				    struct edgetpu_map_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
 	struct edgetpu_map_ioctl ibuf;
 	int ret;
 
@@ -235,63 +280,84 @@ static int etdirect_map_buffer(struct edgetpu_device_group *group,
 
 	trace_edgetpu_map_buffer_start(&ibuf);
 
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	/* to prevent group being released when we perform map/unmap later */
+	group = edgetpu_device_group_get(group);
+	/*
+	 * Don't hold @client->group_lock on purpose since
+	 * 1. We don't care whether @client still belongs to @group.
+	 * 2. get_user_pages_fast called by edgetpu_device_group_map() will hold
+	 *    mm->mmap_sem, we need to prevent our locks being held around it.
+	 */
+	UNLOCK(client);
 	ret = edgetpu_device_group_map(group, &ibuf);
 	if (ret)
-		return ret;
+		goto out;
 
 	if (copy_to_user(argp, &ibuf, sizeof(ibuf))) {
 		edgetpu_device_group_unmap(group, ibuf.die_index,
 					   ibuf.device_address,
 					   EDGETPU_MAP_SKIP_CPU_SYNC);
-		return -EFAULT;
+		ret = -EFAULT;
 	}
 
+out:
+	edgetpu_device_group_put(group);
 	trace_edgetpu_map_buffer_end(&ibuf);
 
-	return 0;
+	return ret;
 }
 
-static int etdirect_unmap_buffer(struct edgetpu_device_group *group,
-				 struct edgetpu_map_ioctl __user *argp)
+static int edgetpu_ioctl_unmap_buffer(struct edgetpu_client *client,
+				      struct edgetpu_map_ioctl __user *argp)
 {
 	struct edgetpu_map_ioctl ibuf;
+	struct edgetpu_device_group *group;
+	int ret;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
 
-	return edgetpu_device_group_unmap(group, ibuf.die_index,
-					  ibuf.device_address, ibuf.flags);
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_device_group_unmap(group, ibuf.die_index,
+					 ibuf.device_address, ibuf.flags);
+	UNLOCK(client);
+	return ret;
 }
 
-static int
-etdirect_allocate_device_buffer_compat(struct edgetpu_device_group *group,
-				struct edgetpu_device_buffer_ioctl __user *argp)
+static int edgetpu_ioctl_allocate_device_buffer_compat(
+	struct edgetpu_client *leader,
+	struct edgetpu_device_buffer_ioctl __user *argp)
 {
 #ifndef EDGETPU_HAS_DEVICE_DRAM
 	return -ENOTTY;
 #else
 	struct edgetpu_device_buffer_ioctl ibuf;
 	struct edgetpu_client *client;
+	struct edgetpu_device_group *group;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
 
+	LOCK_RETURN_IF_NOT_LEADER(leader, group);
 	mutex_lock(&group->lock);
 
 	if (!edgetpu_device_group_is_finalized(group) ||
-			ibuf.die_index >= group->n_clients) {
+	    ibuf.die_index >= group->n_clients) {
 		mutex_unlock(&group->lock);
+		UNLOCK(leader);
 		return -EINVAL;
 	}
 	client = group->members[ibuf.die_index];
 	mutex_unlock(&group->lock);
+	UNLOCK(leader);
 
 	return edgetpu_device_dram_getfd(client, ibuf.size);
 #endif /* EDGETPU_HAS_DEVICE_DRAM */
 }
 
 static int
-etdirect_allocate_device_buffer(struct edgetpu_client *client, u64 size)
+edgetpu_ioctl_allocate_device_buffer(struct edgetpu_client *client, u64 size)
 {
 #ifndef EDGETPU_HAS_DEVICE_DRAM
 	return -ENOTTY;
@@ -300,19 +366,26 @@ etdirect_allocate_device_buffer(struct edgetpu_client *client, u64 size)
 #endif /* EDGETPU_HAS_DEVICE_DRAM */
 }
 
-static int etdirect_sync_buffer(struct edgetpu_device_group *group,
-				struct edgetpu_sync_ioctl __user *argp)
+static int edgetpu_ioctl_sync_buffer(struct edgetpu_client *client,
+				     struct edgetpu_sync_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
+	int ret;
 	struct edgetpu_sync_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	return edgetpu_device_group_sync_buffer(group, &ibuf);
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_device_group_sync_buffer(group, &ibuf);
+	UNLOCK(client);
+	return ret;
 }
 
-static int etdirect_map_dmabuf(struct edgetpu_device_group *group,
-			       struct edgetpu_map_dmabuf_ioctl __user *argp)
+static int
+edgetpu_ioctl_map_dmabuf(struct edgetpu_client *client,
+			 struct edgetpu_map_dmabuf_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
 	struct edgetpu_map_dmabuf_ioctl ibuf;
 	int ret;
 
@@ -321,29 +394,41 @@ static int etdirect_map_dmabuf(struct edgetpu_device_group *group,
 
 	trace_edgetpu_map_dmabuf_start(&ibuf);
 
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	/* to prevent group being released when we perform unmap on fault */
+	group = edgetpu_device_group_get(group);
 	ret = edgetpu_map_dmabuf(group, &ibuf);
+	UNLOCK(client);
 	if (ret)
-		return ret;
+		goto out;
 
 	if (copy_to_user(argp, &ibuf, sizeof(ibuf))) {
 		edgetpu_unmap_dmabuf(group, ibuf.die_index,
 				     ibuf.device_address);
-		return -EFAULT;
+		ret = -EFAULT;
 	}
 
+out:
+	edgetpu_device_group_put(group);
 	trace_edgetpu_map_dmabuf_end(&ibuf);
 
-	return 0;
+	return ret;
 }
 
-static int etdirect_unmap_dmabuf(struct edgetpu_device_group *group,
-				 struct edgetpu_map_dmabuf_ioctl __user *argp)
+static int
+edgetpu_ioctl_unmap_dmabuf(struct edgetpu_client *client,
+			   struct edgetpu_map_dmabuf_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
+	int ret;
 	struct edgetpu_map_dmabuf_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	return edgetpu_unmap_dmabuf(group, ibuf.die_index, ibuf.device_address);
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_unmap_dmabuf(group, ibuf.die_index, ibuf.device_address);
+	UNLOCK(client);
+	return ret;
 }
 
 static int edgetpu_ioctl_sync_fence_create(
@@ -373,35 +458,47 @@ static int edgetpu_ioctl_sync_fence_signal(
 }
 
 static int
-edgetpu_ioctl_map_bulk_dmabuf(struct edgetpu_device_group *group,
+edgetpu_ioctl_map_bulk_dmabuf(struct edgetpu_client *client,
 			      struct edgetpu_map_bulk_dmabuf_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
 	struct edgetpu_map_bulk_dmabuf_ioctl ibuf;
 	int ret;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
 
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	/* to prevent group being released when we perform unmap on fault */
+	group = edgetpu_device_group_get(group);
 	ret = edgetpu_map_bulk_dmabuf(group, &ibuf);
+	UNLOCK(client);
 	if (ret)
-		return ret;
+		goto out;
 
 	if (copy_to_user(argp, &ibuf, sizeof(ibuf))) {
 		edgetpu_unmap_bulk_dmabuf(group, ibuf.device_address);
-		return -EFAULT;
+		ret = -EFAULT;
 	}
-	return 0;
+out:
+	edgetpu_device_group_put(group);
+	return ret;
 }
 
 static int edgetpu_ioctl_unmap_bulk_dmabuf(
-	struct edgetpu_device_group *group,
+	struct edgetpu_client *client,
 	struct edgetpu_map_bulk_dmabuf_ioctl __user *argp)
 {
+	struct edgetpu_device_group *group;
+	int ret;
 	struct edgetpu_map_bulk_dmabuf_ioctl ibuf;
 
 	if (copy_from_user(&ibuf, argp, sizeof(ibuf)))
 		return -EFAULT;
-	return edgetpu_unmap_bulk_dmabuf(group, ibuf.device_address);
+	LOCK_RETURN_IF_NOT_LEADER(client, group);
+	ret = edgetpu_unmap_bulk_dmabuf(group, ibuf.device_address);
+	UNLOCK(client);
+	return ret;
 }
 
 static int edgetpu_ioctl_sync_fence_status(
@@ -430,51 +527,12 @@ static int edgetpu_ioctl_fw_version(struct edgetpu_dev *etdev,
 	return 0;
 }
 
-static bool etdirect_ioctl_check_permissions(struct file *file, uint cmd)
+static bool edgetpu_ioctl_check_permissions(struct file *file, uint cmd)
 {
 	return file->f_mode & FMODE_WRITE;
 }
 
-/*
- * Checks if the state of @client is valid to execute ioctl command @cmd.
- * Caller holds @client->group_lock;
- */
-static bool etdirect_ioctl_check_group(struct edgetpu_client *client, uint cmd)
-{
-	/* @client must not belong to any group */
-	if (cmd == EDGETPU_CREATE_GROUP || cmd == EDGETPU_JOIN_GROUP_COMPAT ||
-	    cmd == EDGETPU_JOIN_GROUP)
-		return !client->group;
-
-	/* Valid for any @client */
-	if (cmd == EDGETPU_SET_PERDIE_EVENTFD ||
-	    cmd == EDGETPU_SET_PERDIE_EVENTFD_COMPAT ||
-	    cmd == EDGETPU_UNSET_PERDIE_EVENT ||
-	    cmd == EDGETPU_UNSET_PERDIE_EVENT_COMPAT ||
-	    cmd == EDGETPU_ALLOCATE_DEVICE_BUFFER ||
-	    cmd == EDGETPU_ALLOCATE_DEVICE_BUFFER_COMPAT_2 ||
-	    cmd == EDGETPU_CREATE_SYNC_FENCE ||
-	    cmd == EDGETPU_SIGNAL_SYNC_FENCE ||
-	    cmd == EDGETPU_SIGNAL_SYNC_FENCE_COMPAT ||
-	    cmd == EDGETPU_SYNC_FENCE_STATUS ||
-	    cmd == EDGETPU_RELEASE_WAKE_LOCK ||
-	    cmd == EDGETPU_ACQUIRE_WAKE_LOCK ||
-	    cmd == EDGETPU_FIRMWARE_VERSION)
-		return true;
-
-	if (!client->group)
-		return false;
-
-	/* Other operations can only be applied on the group lead by @client. */
-	/*
-	 * Note: Though this function already checks the group is not disbanded,
-	 * the callbacks of ioctl still need to check the state of group is
-	 * waiting/finalized with lock to prevent racing.
-	 */
-	return edgetpu_device_group_is_leader(client->group, client);
-}
-
-static int etdirect_ioctl_release_wakelock(struct edgetpu_client *client)
+static int edgetpu_ioctl_release_wakelock(struct edgetpu_client *client)
 {
 	if (!client->etdev->pm)
 		return -ENODEV;
@@ -508,7 +566,7 @@ static int etdirect_ioctl_release_wakelock(struct edgetpu_client *client)
 	return 0;
 }
 
-static int etdirect_ioctl_acquire_wakelock(struct edgetpu_client *client)
+static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 {
 	int ret;
 
@@ -543,83 +601,60 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 	if (!client)
 		return -ENODEV;
 
-	if (!etdirect_ioctl_check_permissions(file, cmd))
+	if (!edgetpu_ioctl_check_permissions(file, cmd))
 		return -EPERM;
 
-	mutex_lock(&client->group_lock);
-	if (!etdirect_ioctl_check_group(client, cmd)) {
-		mutex_unlock(&client->group_lock);
-		return -EINVAL;
-	}
-	/* ioctl commands operating on device group */
 	switch (cmd) {
 	case EDGETPU_MAP_BUFFER:
-		ret = etdirect_map_buffer(client->group, argp);
+		ret = edgetpu_ioctl_map_buffer(client, argp);
 		break;
 	case EDGETPU_UNMAP_BUFFER:
 	case EDGETPU_UNMAP_BUFFER_COMPAT:
-		ret = etdirect_unmap_buffer(client->group, argp);
-		break;
-	case EDGETPU_UNSET_EVENT:
-	case EDGETPU_UNSET_EVENT_COMPAT:
-		edgetpu_group_unset_eventfd(client->group, arg);
-		ret = 0;
-		break;
-	case EDGETPU_FINALIZE_GROUP:
-		ret = edgetpu_device_group_finalize(client->group);
-		break;
-	case EDGETPU_ALLOCATE_DEVICE_BUFFER_COMPAT:
-		ret = etdirect_allocate_device_buffer_compat(client->group,
-							     argp);
-		break;
-	case EDGETPU_SYNC_BUFFER:
-	case EDGETPU_SYNC_BUFFER_COMPAT:
-		ret = etdirect_sync_buffer(client->group, argp);
-		break;
-	case EDGETPU_MAP_DMABUF:
-		ret = etdirect_map_dmabuf(client->group, argp);
-		break;
-	case EDGETPU_UNMAP_DMABUF:
-	case EDGETPU_UNMAP_DMABUF_COMPAT:
-		ret = etdirect_unmap_dmabuf(client->group, argp);
-		break;
-	case EDGETPU_MAP_BULK_DMABUF:
-		ret = edgetpu_ioctl_map_bulk_dmabuf(client->group, argp);
-		break;
-	case EDGETPU_UNMAP_BULK_DMABUF:
-	case EDGETPU_UNMAP_BULK_DMABUF_COMPAT:
-		ret = edgetpu_ioctl_unmap_bulk_dmabuf(client->group, argp);
+		ret = edgetpu_ioctl_unmap_buffer(client, argp);
 		break;
 	case EDGETPU_SET_EVENTFD:
 	case EDGETPU_SET_EVENTFD_COMPAT:
-		ret = etdirect_set_eventfd(client->group, argp);
+		ret = edgetpu_ioctl_set_eventfd(client, argp);
 		break;
-	default:
-		ret = -ENOTTY; /* unknown command */
-	}
-	mutex_unlock(&client->group_lock);
-	if (ret != -ENOTTY)
-		return ret;
-
-	switch (cmd) {
 	case EDGETPU_CREATE_GROUP:
 		ret = edgetpu_ioctl_create_group(client, argp);
 		break;
 	case EDGETPU_JOIN_GROUP:
 	case EDGETPU_JOIN_GROUP_COMPAT:
-		ret = etdirect_join_group(client, (u64)argp);
+		ret = edgetpu_ioctl_join_group(client, (u64)argp);
+		break;
+	case EDGETPU_FINALIZE_GROUP:
+		ret = edgetpu_ioctl_finalize_group(client);
 		break;
 	case EDGETPU_SET_PERDIE_EVENTFD:
 	case EDGETPU_SET_PERDIE_EVENTFD_COMPAT:
-		ret = etdirect_set_perdie_eventfd(client->etdev, argp);
+		ret = edgetpu_ioctl_set_perdie_eventfd(client->etdev, argp);
+		break;
+	case EDGETPU_ALLOCATE_DEVICE_BUFFER_COMPAT:
+		ret = edgetpu_ioctl_allocate_device_buffer_compat(client, argp);
+		break;
+	case EDGETPU_UNSET_EVENT:
+	case EDGETPU_UNSET_EVENT_COMPAT:
+		ret = edgetpu_ioctl_unset_eventfd(client, arg);
 		break;
 	case EDGETPU_UNSET_PERDIE_EVENT:
 	case EDGETPU_UNSET_PERDIE_EVENT_COMPAT:
-		ret = etdirect_unset_perdie_eventfd(client->etdev, arg);
+		ret = edgetpu_ioctl_unset_perdie_eventfd(client->etdev, arg);
+		break;
+	case EDGETPU_SYNC_BUFFER:
+	case EDGETPU_SYNC_BUFFER_COMPAT:
+		ret = edgetpu_ioctl_sync_buffer(client, argp);
+		break;
+	case EDGETPU_MAP_DMABUF:
+		ret = edgetpu_ioctl_map_dmabuf(client, argp);
+		break;
+	case EDGETPU_UNMAP_DMABUF:
+	case EDGETPU_UNMAP_DMABUF_COMPAT:
+		ret = edgetpu_ioctl_unmap_dmabuf(client, argp);
 		break;
 	case EDGETPU_ALLOCATE_DEVICE_BUFFER:
 	case EDGETPU_ALLOCATE_DEVICE_BUFFER_COMPAT_2:
-		ret = etdirect_allocate_device_buffer(client, (u64)argp);
+		ret = edgetpu_ioctl_allocate_device_buffer(client, (u64)argp);
 		break;
 	case EDGETPU_CREATE_SYNC_FENCE:
 		ret = edgetpu_ioctl_sync_fence_create(argp);
@@ -628,14 +663,21 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 	case EDGETPU_SIGNAL_SYNC_FENCE_COMPAT:
 		ret = edgetpu_ioctl_sync_fence_signal(argp);
 		break;
+	case EDGETPU_MAP_BULK_DMABUF:
+		ret = edgetpu_ioctl_map_bulk_dmabuf(client, argp);
+		break;
+	case EDGETPU_UNMAP_BULK_DMABUF:
+	case EDGETPU_UNMAP_BULK_DMABUF_COMPAT:
+		ret = edgetpu_ioctl_unmap_bulk_dmabuf(client, argp);
+		break;
 	case EDGETPU_SYNC_FENCE_STATUS:
 		ret = edgetpu_ioctl_sync_fence_status(argp);
 		break;
 	case EDGETPU_RELEASE_WAKE_LOCK:
-		ret = etdirect_ioctl_release_wakelock(client);
+		ret = edgetpu_ioctl_release_wakelock(client);
 		break;
 	case EDGETPU_ACQUIRE_WAKE_LOCK:
-		ret = etdirect_ioctl_acquire_wakelock(client);
+		ret = edgetpu_ioctl_acquire_wakelock(client);
 		break;
 	case EDGETPU_FIRMWARE_VERSION:
 		ret = edgetpu_ioctl_fw_version(client->etdev, argp);
@@ -647,7 +689,7 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 	return ret;
 }
 
-static long edgetpu_dev_ioctl(struct file *file, uint cmd, ulong arg)
+static long edgetpu_fs_ioctl(struct file *file, uint cmd, ulong arg)
 {
 	if (file->f_op != &edgetpu_fops)
 		return -ENOTTY;
@@ -656,7 +698,7 @@ static long edgetpu_dev_ioctl(struct file *file, uint cmd, ulong arg)
 }
 
 /* Map a region of device/coherent memory. */
-static int etdirect_mmap(struct file *file, struct vm_area_struct *vma)
+static int edgetpu_fs_mmap(struct file *file, struct vm_area_struct *vma)
 {
 	struct edgetpu_client *client = file->private_data;
 
@@ -863,7 +905,7 @@ static const struct file_operations mappings_ops = {
 	.release = single_release,
 };
 
-static void edgetpu_dev_setup_debugfs(struct edgetpu_dev *etdev)
+static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 {
 	etdev->d_entry =
 		debugfs_create_dir(etdev->dev_name, edgetpu_debugfs_dir);
@@ -878,14 +920,14 @@ static void edgetpu_dev_setup_debugfs(struct edgetpu_dev *etdev)
 const struct file_operations edgetpu_fops = {
 	.owner = THIS_MODULE,
 	.llseek = no_llseek,
-	.mmap = etdirect_mmap,
-	.open = edgetpu_dev_open,
-	.release = etdirect_release,
-	.unlocked_ioctl = edgetpu_dev_ioctl,
+	.mmap = edgetpu_fs_mmap,
+	.open = edgetpu_fs_open,
+	.release = edgetpu_fs_release,
+	.unlocked_ioctl = edgetpu_fs_ioctl,
 };
 
 /* Called from edgetpu core to add a new edgetpu device. */
-int edgetpu_dev_add(struct edgetpu_dev *etdev)
+int edgetpu_fs_add(struct edgetpu_dev *etdev)
 {
 	int ret;
 
@@ -911,11 +953,11 @@ int edgetpu_dev_add(struct edgetpu_dev *etdev)
 		return ret;
 	}
 
-	edgetpu_dev_setup_debugfs(etdev);
+	edgetpu_fs_setup_debugfs(etdev);
 	return 0;
 }
 
-void edgetpu_dev_remove(struct edgetpu_dev *etdev)
+void edgetpu_fs_remove(struct edgetpu_dev *etdev)
 {
 	device_destroy(edgetpu_class, etdev->devno);
 	cdev_del(&etdev->cdev);
@@ -943,7 +985,7 @@ static void edgetpu_debugfs_global_setup(void)
 			    &syncfences_ops);
 }
 
-int __init edgetpu_dev_init(void)
+int __init edgetpu_fs_init(void)
 {
 	int ret;
 
@@ -967,14 +1009,14 @@ int __init edgetpu_dev_init(void)
 	return 0;
 }
 
-void __exit edgetpu_dev_exit(void)
+void __exit edgetpu_fs_exit(void)
 {
 	debugfs_remove_recursive(edgetpu_debugfs_dir);
 	unregister_chrdev_region(edgetpu_basedev, EDGETPU_DEV_MAX);
 	class_destroy(edgetpu_class);
 }
 
-struct dentry *edgetpu_dev_debugfs_dir(void)
+struct dentry *edgetpu_fs_debugfs_dir(void)
 {
 	return edgetpu_debugfs_dir;
 }
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index e3e8e20..e2f35f0 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -254,14 +254,14 @@ void edgetpu_unregister_irq(struct edgetpu_dev *etdev, int irq);
 /* Called from core to chip layer when MMU is needed during device init. */
 void edgetpu_setup_mmu(struct edgetpu_dev *etdev);
 
-/* Core -> Device API */
+/* Core -> Device FS API */
 
-int __init edgetpu_dev_init(void);
-void __exit edgetpu_dev_exit(void);
-int edgetpu_dev_add(struct edgetpu_dev *etdev);
-void edgetpu_dev_remove(struct edgetpu_dev *dev);
+int __init edgetpu_fs_init(void);
+void __exit edgetpu_fs_exit(void);
+int edgetpu_fs_add(struct edgetpu_dev *etdev);
+void edgetpu_fs_remove(struct edgetpu_dev *dev);
 /* Get the top-level debugfs directory for the device class */
-struct dentry *edgetpu_dev_debugfs_dir(void);
+struct dentry *edgetpu_fs_debugfs_dir(void);
 
 /* Core/Device -> Chip API */
 
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index ad92fd4..ca8b28c 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -666,19 +666,52 @@ int edgetpu_kci_leave_group(struct edgetpu_kci *kci)
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
-enum edgetpu_fw_flavor edgetpu_kci_fw_flavor(struct edgetpu_kci *kci)
+enum edgetpu_fw_flavor edgetpu_kci_fw_info(
+	struct edgetpu_kci *kci, struct edgetpu_fw_info *fw_info)
 {
+	struct edgetpu_dev *etdev = kci->mailbox->etdev;
 	struct edgetpu_command_element cmd = {
-		.code = KCI_CODE_FIRMWARE_FLAVOR,
+		.code = KCI_CODE_FIRMWARE_INFO,
+		.dma = {
+			.address = 0,
+			.size = 0,
+		},
 	};
+	/* TODO(b/136208139): remove when old fw no longer in use */
+	struct edgetpu_command_element cmd_compat = {
+		.code = KCI_CODE_FIRMWARE_FLAVOR_COMPAT,
+	};
+	dma_addr_t dma_addr;
+	const u32 flags = EDGETPU_MMU_DIE | EDGETPU_MMU_32 | EDGETPU_MMU_HOST;
 	struct edgetpu_kci_response_element resp;
 	enum edgetpu_fw_flavor flavor = FW_FLAVOR_UNKNOWN;
 	int kciret;
 
+	dma_addr = dma_map_single(etdev->dev, fw_info, sizeof(*fw_info),
+				  DMA_FROM_DEVICE);
+	/* If any map failure still try handshake without full fw_info */
+	if (dma_mapping_error(etdev->dev, dma_addr)) {
+		etdev_warn(etdev, "%s: failed to DMA map fw info buffer",
+			  __func__);
+	} else {
+		cmd.dma.address =
+			edgetpu_mmu_tpu_map(etdev, dma_addr, sizeof(*fw_info),
+					    DMA_FROM_DEVICE,
+					    EDGETPU_CONTEXT_KCI, flags);
+		if (!cmd.dma.address)
+			etdev_warn(etdev,
+				   "%s: failed to map fw info buffer to TPU",
+				   __func__);
+		else
+			cmd.dma.size = sizeof(*fw_info);
+	}
+
 	kciret = edgetpu_kci_send_cmd_return_resp(kci, &cmd, &resp);
+	if (kciret == KCI_ERROR_UNIMPLEMENTED)
+		kciret = edgetpu_kci_send_cmd_return_resp(kci, &cmd_compat,
+							  &resp);
 	if (kciret == KCI_ERROR_UNIMPLEMENTED) {
-		etdev_dbg(kci->mailbox->etdev,
-			  "old firmware does not report flavor\n");
+		etdev_dbg(etdev, "old firmware does not report flavor\n");
 	} else if (kciret == KCI_ERROR_OK) {
 		switch (resp.retval) {
 		case FW_FLAVOR_BL1:
@@ -688,19 +721,23 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_flavor(struct edgetpu_kci *kci)
 			flavor = resp.retval;
 			break;
 		default:
-			etdev_dbg(kci->mailbox->etdev,
-				  "unrecognized fw flavor 0x%x\n",
+			etdev_dbg(etdev, "unrecognized fw flavor 0x%x\n",
 				  resp.retval);
 		}
 	} else {
-		etdev_dbg(kci->mailbox->etdev,
-			  "firmware flavor query returns %d\n", kciret);
+		etdev_dbg(etdev, "firmware flavor query returns %d\n", kciret);
 		if (kciret < 0)
 			flavor = kciret;
 		else
 			flavor = -EIO;
 	}
 
+	if (cmd.dma.address)
+		edgetpu_mmu_tpu_unmap(etdev, cmd.dma.address,
+				      sizeof(*fw_info), EDGETPU_CONTEXT_KCI);
+	if (!dma_mapping_error(etdev->dev, dma_addr))
+		dma_unmap_single(etdev->dev, dma_addr, sizeof(*fw_info),
+				 DMA_FROM_DEVICE);
 	return flavor;
 }
 
@@ -737,3 +774,31 @@ int edgetpu_kci_shutdown(struct edgetpu_kci *kci)
 		return -ENODEV;
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
+
+int edgetpu_kci_open_device(struct edgetpu_kci *kci, u8 mailbox_id)
+{
+	struct edgetpu_command_element cmd = {
+		.code = KCI_CODE_OPEN_DEVICE,
+		.dma = {
+			.flags = mailbox_id,
+		},
+	};
+
+	if (!kci)
+		return -ENODEV;
+	return edgetpu_kci_send_cmd(kci, &cmd);
+}
+
+int edgetpu_kci_close_device(struct edgetpu_kci *kci, u8 mailbox_id)
+{
+	struct edgetpu_command_element cmd = {
+		.code = KCI_CODE_CLOSE_DEVICE,
+		.dma = {
+			.flags = mailbox_id,
+		},
+	};
+
+	if (!kci)
+		return -ENODEV;
+	return edgetpu_kci_send_cmd(kci, &cmd);
+}
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index 7db571b..83b03a7 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -16,6 +16,7 @@
 #include <linux/types.h>
 #include <linux/wait.h>
 
+#include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-mailbox.h"
 
@@ -87,8 +88,12 @@ enum edgetpu_kci_code {
 	KCI_CODE_JOIN_GROUP = 3,
 	KCI_CODE_LEAVE_GROUP = 4,
 	KCI_CODE_MAP_TRACE_BUFFER = 5,
-	KCI_CODE_FIRMWARE_FLAVOR = 6,
+	/* TODO(b/136208139): remove when old fw no longer in use */
+	KCI_CODE_FIRMWARE_FLAVOR_COMPAT = 6,
 	KCI_CODE_SHUTDOWN = 7,
+	KCI_CODE_OPEN_DEVICE = 9,
+	KCI_CODE_CLOSE_DEVICE = 10,
+	KCI_CODE_FIRMWARE_INFO = 11,
 };
 
 /*
@@ -115,20 +120,6 @@ enum edgetpu_kci_error {
 	KCI_ERROR_UNAUTHENTICATED = 16,
 };
 
-/* Firmware flavors returned via KCI from firmware image. */
-enum edgetpu_fw_flavor {
-	/* used by host when cannot determine the flavor */
-	FW_FLAVOR_UNKNOWN = 0,
-	/* second-stage bootloader */
-	FW_FLAVOR_BL1 = 1,
-	/* systest app image */
-	FW_FLAVOR_SYSTEST = 2,
-	/* default production app image from DarwiNN team */
-	FW_FLAVOR_PROD_DEFAULT = 3,
-	/* custom image produced by other teams */
-	FW_FLAVOR_CUSTOM = 4,
-};
-
 struct edgetpu_kci_wait_list {
 	struct list_head list;
 	/*
@@ -239,14 +230,18 @@ int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
 int edgetpu_kci_ack(struct edgetpu_kci *kci);
 
 /*
- * Sends a FIRMWARE_FLAVOR command and expects a response indicating what
- * edgetpu_fw_flavor type is running.  Also serves as an initial handshake
- * with firmware at load time.
+ * Sends a FIRMWARE_INFO command and expects a response with a
+ * edgetpu_fw_info struct filled out, including what firmware type is running,
+ * along with build CL and time.
+ * Also serves as an initial handshake with firmware at load time.
+ *
+ * @fw_info: a struct edgetpu_fw_info to be filled out by fw
  *
  * Returns >=0 edgetpu_fw_flavor when response received from firmware,
  *         <0 on error communicating with firmware (typically -ETIMEDOUT).
  */
-enum edgetpu_fw_flavor edgetpu_kci_fw_flavor(struct edgetpu_kci *kci);
+enum edgetpu_fw_flavor edgetpu_kci_fw_info(
+	struct edgetpu_kci *kci, struct edgetpu_fw_info *fw_info);
 
 /*
  * Sends the "Map Log Buffer" command and waits for remote response.
@@ -281,4 +276,10 @@ void edgetpu_kci_mappings_show(struct edgetpu_dev *etdev, struct seq_file *s);
 /* Send shutdown request to firmware */
 int edgetpu_kci_shutdown(struct edgetpu_kci *kci);
 
+/* Inform the firmware to prepare to serve the VII with @mailbox_id. */
+int edgetpu_kci_open_device(struct edgetpu_kci *kci, u8 mailbox_id);
+
+/* Inform the firmware the VII with @mailbox_id is closed. */
+int edgetpu_kci_close_device(struct edgetpu_kci *kci, u8 mailbox_id);
+
 #endif /* __EDGETPU_KCI_H__ */
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index a62da18..86139a8 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -10,14 +10,19 @@
 #include <linux/slab.h>
 
 #include "edgetpu-config.h"
-#include "edgetpu-device-group.h"
-#include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 #include "edgetpu-pm.h"
 #include "edgetpu-sw-watchdog.h"
 
+#if IS_ENABLED(CONFIG_EDGETPU_TEST)
+#include "unittests/factory/fake-edgetpu-firmware.h"
+#define SIM_PCHANNEL(etdev) fake_edgetpu_firmware_sim_pchannel(etdev)
+#else
+#define SIM_PCHANNEL(...)
+#endif
+
 struct edgetpu_pm_private {
 	const struct edgetpu_pm_handlers *handlers;
 	struct mutex lock;
@@ -123,23 +128,26 @@ void edgetpu_pm_destroy(struct edgetpu_dev *etdev)
 	etdev->pm = NULL;
 }
 
-void edgetpu_pm_shutdown(struct edgetpu_dev *etdev)
+void edgetpu_pm_shutdown(struct edgetpu_dev *etdev, bool force)
 {
 	struct edgetpu_pm *etpm = etdev->pm;
 
 	if (!etpm)
 		return;
 	mutex_lock(&etpm->p->lock);
-	if (etdev->firmware)
-		edgetpu_firmware_lock(etdev);
+
+	/* someone is using the device */
 	if (etpm->p->power_up_count) {
-		etdev_warn(etdev, "Leaving %d clients behind!\n",
-			   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);
-	if (etdev->firmware)
-		edgetpu_firmware_unlock(etdev);
+unlock:
 	mutex_unlock(&etpm->p->lock);
 }
 
@@ -153,6 +161,14 @@ bool edgetpu_is_powered(struct edgetpu_dev *etdev)
 	return etpm->p->power_up_count;
 }
 
+#define etdev_poll_power_state(etdev, val, cond)                               \
+	({                                                                     \
+		SIM_PCHANNEL(etdev);                                           \
+		readl_relaxed_poll_timeout(                                    \
+			etdev->regs.mem + EDGETPU_REG_POWER_CONTROL, val,      \
+			cond, 1, EDGETPU_PCHANNEL_STATE_CHANGE_TIMEOUT);       \
+	})
+
 static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 {
 	int ret;
@@ -165,10 +181,7 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 	if (val & PREQ) {
 		edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL,
 				     val & ~(PREQ));
-		ret = readl_relaxed_poll_timeout(
-			etdev->regs.mem + EDGETPU_REG_POWER_CONTROL, val,
-			(val & PACCEPT) == 0, 1,
-			EDGETPU_PCHANNEL_STATE_CHANGE_TIMEOUT);
+		ret = etdev_poll_power_state(etdev, val, (val & PACCEPT) == 0);
 		if (ret) {
 			etdev_err(etdev, "p-channel request timeout\n");
 			return ret;
@@ -176,16 +189,15 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 	}
 	/* Phase 2: Request state */
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL, state | PREQ);
+	SIM_PCHANNEL(etdev);
 
 	/* don't wait for state accept if STATE RUN */
 	if (state == STATE_RUN)
 		return 0;
 
 	/* Phase 3: R52 acknowledgment */
-	ret = readl_relaxed_poll_timeout(
-		etdev->regs.mem + EDGETPU_REG_POWER_CONTROL, val,
-		(val & PACCEPT) || (val & PDENY), 1,
-		EDGETPU_PCHANNEL_STATE_CHANGE_TIMEOUT);
+	ret = etdev_poll_power_state(etdev, val,
+				     (val & PACCEPT) || (val & PDENY));
 	if (val & PDENY) {
 		edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL,
 				     val & !state);
@@ -198,10 +210,8 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 	}
 	/* Phase 4. Drive PREQ to 0 */
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL, val & ~(PREQ));
-	ret = readl_relaxed_poll_timeout(
-		etdev->regs.mem + EDGETPU_REG_POWER_CONTROL, val,
-		((val & PACCEPT) == 0) && ((val & PDENY) == 0), 1,
-		EDGETPU_PCHANNEL_STATE_CHANGE_TIMEOUT);
+	ret = etdev_poll_power_state(
+		etdev, val, ((val & PACCEPT) == 0) && ((val & PDENY) == 0));
 
 	return deny ? -EACCES : ret;
 }
@@ -226,10 +236,7 @@ int edgetpu_pchannel_power_down(struct edgetpu_dev *etdev, bool wait_on_pactive)
 		msleep(200);
 	else
 		/* wait for PACTIVE[1] goes low. */
-		ret = readl_relaxed_poll_timeout(
-			etdev->regs.mem + EDGETPU_REG_POWER_CONTROL, val,
-			(val & PACTIVE) == 0, 1,
-			EDGETPU_PCHANNEL_STATE_CHANGE_TIMEOUT);
+		ret = etdev_poll_power_state(etdev, val, (val & PACTIVE) == 0);
 	if (ret)
 		return ret;
 	ret = pchannel_state_change_request(etdev, STATE_SHUTDOWN);
diff --git a/drivers/edgetpu/edgetpu-pm.h b/drivers/edgetpu/edgetpu-pm.h
index b7c9b5b..abe042a 100644
--- a/drivers/edgetpu/edgetpu-pm.h
+++ b/drivers/edgetpu/edgetpu-pm.h
@@ -38,7 +38,7 @@ struct edgetpu_pm {
  * 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-direct and edgetpu-firmware
+ * 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.
@@ -62,10 +62,12 @@ int edgetpu_pm_create(struct edgetpu_dev *etdev,
 void edgetpu_pm_destroy(struct edgetpu_dev *etdev);
 
 /*
- * Ensure device is shut down, should be a no-op unless a client was left open
- * while the device is being removed.
+ * 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);
+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);
diff --git a/drivers/edgetpu/edgetpu-sw-watchdog.c b/drivers/edgetpu/edgetpu-sw-watchdog.c
index 4cb96e3..397c021 100644
--- a/drivers/edgetpu/edgetpu-sw-watchdog.c
+++ b/drivers/edgetpu/edgetpu-sw-watchdog.c
@@ -4,6 +4,7 @@
  *
  * Copyright (C) 2020 Google, Inc.
  */
+#include <linux/module.h>
 #include <linux/slab.h>
 #include <linux/workqueue.h>
 
@@ -11,6 +12,9 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-sw-watchdog.h"
 
+static bool wdt_disable;
+module_param(wdt_disable, bool, 0660);
+
 /* Worker to execute action callback handler on watchdog bite. */
 static void sw_wdt_handler_work(struct work_struct *work)
 {
@@ -60,6 +64,7 @@ int edgetpu_sw_wdt_create(struct edgetpu_dev *etdev, unsigned long hrtbeat_ms)
 	etdev_sw_wdt->hrtbeat_jiffs = msecs_to_jiffies(hrtbeat_ms);
 	INIT_DELAYED_WORK(&etdev_sw_wdt->dwork, sw_wdt_work);
 	INIT_WORK(&etdev_sw_wdt->et_action_work.work, sw_wdt_handler_work);
+	etdev_sw_wdt->is_wdt_disabled = wdt_disable;
 	etdev->etdev_sw_wdt = etdev_sw_wdt;
 	return 0;
 }
@@ -72,6 +77,10 @@ int edgetpu_sw_wdt_start(struct edgetpu_dev *etdev)
 		return -EINVAL;
 	if (!etdev_sw_wdt->et_action_work.edgetpu_sw_wdt_handler)
 		etdev_err(etdev, "sw wdt handler not set\n");
+	if (etdev_sw_wdt->is_wdt_disabled) {
+		etdev_dbg(etdev, "sw wdt disabled by module param");
+		return 0;
+	}
 	etdev_dbg(etdev, "sw wdt: started\n");
 	schedule_delayed_work(&etdev_sw_wdt->dwork,
 			      etdev_sw_wdt->hrtbeat_jiffs);
diff --git a/drivers/edgetpu/edgetpu-sw-watchdog.h b/drivers/edgetpu/edgetpu-sw-watchdog.h
index 931cc08..c278912 100644
--- a/drivers/edgetpu/edgetpu-sw-watchdog.h
+++ b/drivers/edgetpu/edgetpu-sw-watchdog.h
@@ -30,6 +30,8 @@ struct edgetpu_sw_wdt {
 	unsigned long hrtbeat_jiffs;
 	/* work information for watchdog bite. */
 	struct edgetpu_sw_wdt_action_work et_action_work;
+	/* flag to mark that watchdog is disabled. */
+	bool is_wdt_disabled;
 };
 
 int edgetpu_sw_wdt_create(struct edgetpu_dev *etdev, unsigned long hrtbeat_ms);
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index 04daf4c..928c3c9 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -231,18 +231,21 @@ static void edgetpu_fw_log(struct edgetpu_dev *etdev,
 		copy_with_wrap(header, buffer, entry.length, queue_size, start);
 		buffer[entry.length] = 0;
 
+		if (entry.code > EDGETPU_FW_DMESG_LOG_LEVEL)
+			continue;
+
 		switch (entry.code) {
-		case 2:
-		case 1:
+		case EDGETPU_FW_LOG_LEVEL_VERBOSE:
+		case EDGETPU_FW_LOG_LEVEL_DEBUG:
 			etdev_dbg_ratelimited(etdev, "%s", buffer);
 			break;
-		case -1:
+		case EDGETPU_FW_LOG_LEVEL_WARN:
 			etdev_warn_ratelimited(etdev, "%s", buffer);
 			break;
-		case -2:
+		case EDGETPU_FW_LOG_LEVEL_ERROR:
 			etdev_err_ratelimited(etdev, "%s", buffer);
 			break;
-		case 0:
+		case EDGETPU_FW_LOG_LEVEL_INFO:
 		default:
 			etdev_info_ratelimited(etdev, "%s", buffer);
 			break;
diff --git a/drivers/edgetpu/edgetpu-telemetry.h b/drivers/edgetpu/edgetpu-telemetry.h
index 3005041..879fcb6 100644
--- a/drivers/edgetpu/edgetpu-telemetry.h
+++ b/drivers/edgetpu/edgetpu-telemetry.h
@@ -16,6 +16,15 @@
 #include "edgetpu-internal.h"
 #include "edgetpu-kci.h"
 
+/* Log level codes used by edgetpu firmware */
+#define EDGETPU_FW_LOG_LEVEL_VERBOSE (2)
+#define EDGETPU_FW_LOG_LEVEL_DEBUG (1)
+#define EDGETPU_FW_LOG_LEVEL_INFO (0)
+#define EDGETPU_FW_LOG_LEVEL_WARN (-1)
+#define EDGETPU_FW_LOG_LEVEL_ERROR (-2)
+
+#define EDGETPU_FW_DMESG_LOG_LEVEL (EDGETPU_FW_LOG_LEVEL_ERROR)
+
 #define EDGETPU_TELEMETRY_BUFFER_SIZE (16 * 4096)
 /* assumes buffer size is power of 2 */
 #define EDGETPU_TELEMETRY_WRAP_BIT EDGETPU_TELEMETRY_BUFFER_SIZE