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

gcip: Add gcip-telemetry
Bug: 239374826
edgetpu: Use devm_kstrdup
edgetpu: Adopt GCIP telemetry functions
Bug: 239374826 (repeat)
edgetpu: Adopt GCIP telemetry struct
Bug: 239374826 (repeat)
edgetpu: Adopt GCIP telemetry enum and define
Bug: 239374826 (repeat)
gcip: Add gcip-telemetry
Bug: 239374826 (repeat)
edgetpu: move fw context set call to mobile common
edgetpu: remove obsolete define for number of tiles
edgetpu: remove obsolete status register dump defines
edgetpu: make edgetpu_fops static
edgetpu: add comment on usage of "get clock rate" flags param
edgetpu: move mobile chip firmware API calls to common
edgetpu: remove unused platform power callbacks
edgetpu: gsx01: update dummy bcl header
edgetpu: rio: move BCL code from common to rio PM
edgetpu: remove unused firmware handler teardown_buffer
edgetpu: Kbuild: Append internal headers after EXTRA_CFLAGS
Bug: 245705573

Signed-off-by: Zuma copybara merger <zuma-automerger@google.com>
GitOrigin-RevId: 4e74b17fc3f777acad0131c89416803f0a022b70
Change-Id: I1072918998fbfd45f67ec8c6469d2640bac9913e

22 files changed, 616 insertions, 555 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 85e400b..913899b 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -1,7 +1,7 @@
 obj-m	+= rio.o
 # Use the absolute path of this Makefile to get the source directory.
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
-ccflags-y	+= -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -I$(CURRENT_DIR)/include \
+ccflags-y	+= $(EXTRA_CFLAGS) -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -I$(CURRENT_DIR)/include \
 		   -DCONFIG_GCIP=1 -I$(CURRENT_DIR)/gcip-kernel-driver/include
 GIT_PATH=$(CURRENT_DIR)/../../
 ifeq ($(shell git --git-dir=$(GIT_PATH)/.git rev-parse --is-inside-work-tree),true)
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 43ef85a..c501d7d 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -206,11 +206,11 @@ static void edgetpu_vma_open(struct vm_area_struct *vma)
 	/* handle telemetry types */
 	switch (type) {
 	case VMA_LOG:
-		edgetpu_telemetry_inc_mmap_count(etdev, EDGETPU_TELEMETRY_LOG,
+		edgetpu_telemetry_inc_mmap_count(etdev, GCIP_TELEMETRY_LOG,
 						 VMA_DATA_GET(pvt->flag));
 		break;
 	case VMA_TRACE:
-		edgetpu_telemetry_inc_mmap_count(etdev, EDGETPU_TELEMETRY_TRACE,
+		edgetpu_telemetry_inc_mmap_count(etdev, GCIP_TELEMETRY_TRACE,
 						 VMA_DATA_GET(pvt->flag));
 		break;
 	default:
@@ -233,11 +233,11 @@ static void edgetpu_vma_close(struct vm_area_struct *vma)
 	/* handle telemetry types */
 	switch (type) {
 	case VMA_LOG:
-		edgetpu_telemetry_dec_mmap_count(etdev, EDGETPU_TELEMETRY_LOG,
+		edgetpu_telemetry_dec_mmap_count(etdev, GCIP_TELEMETRY_LOG,
 						 VMA_DATA_GET(pvt->flag));
 		break;
 	case VMA_TRACE:
-		edgetpu_telemetry_dec_mmap_count(etdev, EDGETPU_TELEMETRY_TRACE,
+		edgetpu_telemetry_dec_mmap_count(etdev, GCIP_TELEMETRY_TRACE,
 						 VMA_DATA_GET(pvt->flag));
 		break;
 	default:
@@ -299,12 +299,12 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 
 	/* Allow mapping log and telemetry buffers without a group */
 	if (type == VMA_LOG) {
-		ret = edgetpu_mmap_telemetry_buffer(client->etdev, EDGETPU_TELEMETRY_LOG, vma,
+		ret = edgetpu_mmap_telemetry_buffer(client->etdev, GCIP_TELEMETRY_LOG, vma,
 						    VMA_DATA_GET(flag));
 		goto out_set_op;
 	}
 	if (type == VMA_TRACE) {
-		ret = edgetpu_mmap_telemetry_buffer(client->etdev, EDGETPU_TELEMETRY_TRACE, vma,
+		ret = edgetpu_mmap_telemetry_buffer(client->etdev, GCIP_TELEMETRY_TRACE, vma,
 						    VMA_DATA_GET(flag));
 		goto out_set_op;
 	}
@@ -598,10 +598,10 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	/* Clean up all the per die event fds registered by the client */
 	if (client->perdie_events &
 	    1 << perdie_event_id_to_num(EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE))
-		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_LOG);
+		edgetpu_telemetry_unset_event(etdev, GCIP_TELEMETRY_LOG);
 	if (client->perdie_events &
 	    1 << perdie_event_id_to_num(EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE))
-		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_TRACE);
+		edgetpu_telemetry_unset_event(etdev, GCIP_TELEMETRY_TRACE);
 
 	edgetpu_client_put(client);
 }
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index d2756ef..fc26c2f 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -49,6 +49,49 @@ void *edgetpu_firmware_get_data(struct edgetpu_firmware *et_fw)
 	return et_fw->p->data;
 }
 
+/* Request firmware and copy to carveout. */
+static int edgetpu_firmware_carveout_load_locked(struct edgetpu_firmware *et_fw,
+						 struct edgetpu_firmware_desc *fw_desc,
+						 const char *name)
+{
+	int ret;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct device *dev = etdev->dev;
+	const struct firmware *fw;
+	size_t aligned_size;
+
+	ret = request_firmware(&fw, name, dev);
+	if (ret) {
+		etdev_err(etdev, "request firmware '%s' failed: %d\n", name, ret);
+		return ret;
+	}
+
+	aligned_size = ALIGN(fw->size, fw_desc->buf.used_size_align);
+	if (aligned_size > fw_desc->buf.alloc_size) {
+		etdev_err(etdev,
+			  "firmware buffer too small: alloc size=%#zx, required size=%#zx\n",
+			  fw_desc->buf.alloc_size, aligned_size);
+		ret = -ENOSPC;
+		goto out_release_firmware;
+	}
+
+	memcpy(fw_desc->buf.vaddr, fw->data, fw->size);
+	fw_desc->buf.used_size = aligned_size;
+	/* May return NULL on out of memory, driver must handle properly */
+	fw_desc->buf.name = devm_kstrdup(dev, name, GFP_KERNEL);
+
+out_release_firmware:
+	release_firmware(fw);
+	return ret;
+}
+
+static void edgetpu_firmware_carveout_unload_locked(struct edgetpu_firmware *et_fw,
+						    struct edgetpu_firmware_desc *fw_desc)
+{
+	fw_desc->buf.name = NULL;
+	fw_desc->buf.used_size = 0;
+}
+
 static int edgetpu_firmware_load_locked(
 		struct edgetpu_firmware *et_fw,
 		struct edgetpu_firmware_desc *fw_desc, const char *name,
@@ -69,11 +112,9 @@ static int edgetpu_firmware_load_locked(
 		}
 	}
 
-	ret = edgetpu_firmware_chip_load_locked(et_fw, fw_desc, name);
-	if (ret) {
-		etdev_err(etdev, "firmware request failed: %d\n", ret);
+	ret = edgetpu_firmware_carveout_load_locked(et_fw, fw_desc, name);
+	if (ret)
 		goto out_free_buffer;
-	}
 
 	if (chip_fw->setup_buffer) {
 		ret = chip_fw->setup_buffer(et_fw, &fw_desc->buf);
@@ -87,7 +128,7 @@ static int edgetpu_firmware_load_locked(
 	return 0;
 
 out_unload_locked:
-	edgetpu_firmware_chip_unload_locked(et_fw, fw_desc);
+	edgetpu_firmware_carveout_unload_locked(et_fw, fw_desc);
 out_free_buffer:
 	if (chip_fw->free_buffer)
 		chip_fw->free_buffer(et_fw, &fw_desc->buf);
@@ -100,7 +141,7 @@ static void edgetpu_firmware_unload_locked(
 {
 	const struct edgetpu_firmware_chip_data *chip_fw = et_fw->p->chip_fw;
 
-	edgetpu_firmware_chip_unload_locked(et_fw, fw_desc);
+	edgetpu_firmware_carveout_unload_locked(et_fw, fw_desc);
 	/*
 	 * Platform specific implementation for freeing allocated buffer.
 	 */
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index 6660540..f8eb6ac 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -121,9 +121,6 @@ struct edgetpu_firmware_chip_data {
 	 */
 	int (*setup_buffer)(struct edgetpu_firmware *et_fw,
 			    struct edgetpu_firmware_buffer *fw_buf);
-	/* Release the resources previously allocated by setup_buffer(). */
-	void (*teardown_buffer)(struct edgetpu_firmware *et_fw,
-				struct edgetpu_firmware_buffer *fw_buf);
 	/*
 	 * Platform-specific handling after firmware loaded, before running
 	 * the firmware, such as validating the firmware or resetting the
@@ -140,17 +137,6 @@ struct edgetpu_firmware_chip_data {
 };
 
 /*
- * Chip-dependent (actually chip family dependent) calls for loading/unloading
- * firmware images.  Used by the common firmware layer.
- */
-int edgetpu_firmware_chip_load_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc, const char *name);
-void edgetpu_firmware_chip_unload_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc);
-
-/*
  * Returns the chip-specific IOVA where the firmware is mapped.
  *
  * Debug purpose only.
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 47ef902..39730ec 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -72,11 +72,6 @@ static struct dentry *edgetpu_debugfs_dir;
 		}                                                              \
 	} while (0)
 
-bool is_edgetpu_file(struct file *file)
-{
-	return file->f_op == &edgetpu_fops;
-}
-
 int edgetpu_open(struct edgetpu_dev_iface *etiface, struct file *file)
 {
 	struct edgetpu_client *client;
@@ -178,11 +173,9 @@ edgetpu_ioctl_set_perdie_eventfd(struct edgetpu_client *client,
 
 	switch (eventreg.event_id) {
 	case EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE:
-		return edgetpu_telemetry_set_event(etdev, EDGETPU_TELEMETRY_LOG,
-						   eventreg.eventfd);
+		return edgetpu_telemetry_set_event(etdev, GCIP_TELEMETRY_LOG, eventreg.eventfd);
 	case EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE:
-		return edgetpu_telemetry_set_event(
-			etdev, EDGETPU_TELEMETRY_TRACE, eventreg.eventfd);
+		return edgetpu_telemetry_set_event(etdev, GCIP_TELEMETRY_TRACE, eventreg.eventfd);
 	default:
 		return -EINVAL;
 	}
@@ -199,10 +192,10 @@ static int edgetpu_ioctl_unset_perdie_eventfd(struct edgetpu_client *client,
 
 	switch (event_id) {
 	case EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE:
-		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_LOG);
+		edgetpu_telemetry_unset_event(etdev, GCIP_TELEMETRY_LOG);
 		break;
 	case EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE:
-		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_TRACE);
+		edgetpu_telemetry_unset_event(etdev, GCIP_TELEMETRY_TRACE);
 		break;
 	default:
 		return -EINVAL;
@@ -960,7 +953,7 @@ static const struct attribute_group edgetpu_attr_group = {
 	.attrs = edgetpu_dev_attrs,
 };
 
-const struct file_operations edgetpu_fops = {
+static const struct file_operations edgetpu_fops = {
 	.owner = THIS_MODULE,
 	.llseek = no_llseek,
 	.mmap = edgetpu_fs_mmap,
@@ -969,6 +962,11 @@ const struct file_operations edgetpu_fops = {
 	.unlocked_ioctl = edgetpu_fs_ioctl,
 };
 
+bool is_edgetpu_file(struct file *file)
+{
+	return file->f_op == &edgetpu_fops;
+}
+
 static int edgeptu_fs_add_interface(struct edgetpu_dev *etdev, struct edgetpu_dev_iface *etiface,
 				    const struct edgetpu_iface_params *etiparams)
 {
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 5dc4723..7c7a433 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -55,9 +55,6 @@
 #define etdev_warn_once(etdev, fmt, ...)                                       \
 	dev_warn_once(get_dev_for_logging(etdev), fmt, ##__VA_ARGS__)
 
-/* The number of TPU tiles in an edgetpu chip */
-#define EDGETPU_NTILES	16
-
 /*
  * Common-layer context IDs for non-secure TPU access, translated to chip-
  * specific values in the mmu driver.
@@ -242,18 +239,6 @@ enum edgetpu_fw_crash_type {
 	EDGETPU_FW_CRASH_UNRECOV_FAULT = 4,
 };
 
-extern const struct file_operations edgetpu_fops;
-
-/* Status regs dump. */
-struct edgetpu_dumpregs_range {
-	u32 firstreg;
-	u32 lastreg;
-};
-extern struct edgetpu_dumpregs_range edgetpu_chip_statusregs_ranges[];
-extern int edgetpu_chip_statusregs_nranges;
-extern struct edgetpu_dumpregs_range edgetpu_chip_tile_statusregs_ranges[];
-extern int edgetpu_chip_tile_statusregs_nranges;
-
 static inline const char *edgetpu_dma_dir_rw_s(enum dma_data_direction dir)
 {
 	static const char *tbl[4] = { "rw", "r", "w", "?" };
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 2caaee4..644e510 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -15,6 +15,8 @@
 #include <linux/slab.h>
 #include <linux/string.h> /* memcpy */
 
+#include <gcip/gcip-telemetry.h>
+
 #include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -347,30 +349,30 @@ static int edgetpu_kci_send_cmd_with_data(struct edgetpu_kci *etkci,
 	return ret;
 }
 
-int edgetpu_kci_map_log_buffer(struct edgetpu_kci *etkci, tpu_addr_t tpu_addr, u32 size)
+int edgetpu_kci_map_log_buffer(struct gcip_telemetry_kci_args *args)
 {
 	struct gcip_kci_command_element cmd = {
 		.code = GCIP_KCI_CODE_MAP_LOG_BUFFER,
 		.dma = {
-			.address = tpu_addr,
-			.size = size,
+			.address = args->addr,
+			.size = args->size,
 		},
 	};
 
-	return gcip_kci_send_cmd(etkci->kci, &cmd);
+	return gcip_kci_send_cmd(args->kci, &cmd);
 }
 
-int edgetpu_kci_map_trace_buffer(struct edgetpu_kci *etkci, tpu_addr_t tpu_addr, u32 size)
+int edgetpu_kci_map_trace_buffer(struct gcip_telemetry_kci_args *args)
 {
 	struct gcip_kci_command_element cmd = {
 		.code = GCIP_KCI_CODE_MAP_TRACE_BUFFER,
 		.dma = {
-			.address = tpu_addr,
-			.size = size,
+			.address = args->addr,
+			.size = args->size,
 		},
 	};
 
-	return gcip_kci_send_cmd(etkci->kci, &cmd);
+	return gcip_kci_send_cmd(args->kci, &cmd);
 }
 
 enum gcip_fw_flavor edgetpu_kci_fw_info(struct edgetpu_kci *etkci, struct gcip_fw_info *fw_info)
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index 4a4d6c8..0368be6 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -132,19 +132,21 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev);
  */
 int edgetpu_kci_update_usage_locked(struct edgetpu_dev *etdev);
 
+struct gcip_telemetry_kci_args;
+
 /*
  * Sends the "Map Log Buffer" command and waits for remote response.
  *
  * Returns the code of response, or a negative errno on error.
  */
-int edgetpu_kci_map_log_buffer(struct edgetpu_kci *etkci, tpu_addr_t tpu_addr, u32 size);
+int edgetpu_kci_map_log_buffer(struct gcip_telemetry_kci_args *args);
 
 /*
  * Sends the "Map Trace Buffer" command and waits for remote response.
  *
  * Returns the code of response, or a negative errno on error.
  */
-int edgetpu_kci_map_trace_buffer(struct edgetpu_kci *etkci, tpu_addr_t tpu_addr, u32 size);
+int edgetpu_kci_map_trace_buffer(struct gcip_telemetry_kci_args *args);
 
 /* debugfs mappings dump */
 void edgetpu_kci_mappings_show(struct edgetpu_dev *etdev, struct seq_file *s);
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index f74e9b1..c56990a 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -32,11 +32,11 @@
 	 EDGETPU_NUM_CORES)
 
 static void get_telemetry_mem(struct edgetpu_mobile_platform_dev *etmdev,
-			      enum edgetpu_telemetry_type type, struct edgetpu_coherent_mem *mem)
+			      enum gcip_telemetry_type type, struct edgetpu_coherent_mem *mem)
 {
-	int i, offset = type == EDGETPU_TELEMETRY_TRACE ? EDGETPU_TELEMETRY_LOG_BUFFER_SIZE : 0;
-	const size_t size = type == EDGETPU_TELEMETRY_LOG ? EDGETPU_TELEMETRY_LOG_BUFFER_SIZE :
-								EDGETPU_TELEMETRY_TRACE_BUFFER_SIZE;
+	int i, offset = type == GCIP_TELEMETRY_TRACE ? EDGETPU_TELEMETRY_LOG_BUFFER_SIZE : 0;
+	const size_t size = type == GCIP_TELEMETRY_LOG ? EDGETPU_TELEMETRY_LOG_BUFFER_SIZE :
+							 EDGETPU_TELEMETRY_TRACE_BUFFER_SIZE;
 
 	for (i = 0; i < etmdev->edgetpu_dev.num_cores; i++) {
 		mem[i].vaddr = etmdev->shared_mem_vaddr + offset;
@@ -50,8 +50,8 @@ static void get_telemetry_mem(struct edgetpu_mobile_platform_dev *etmdev,
 
 static void edgetpu_mobile_get_telemetry_mem(struct edgetpu_mobile_platform_dev *etmdev)
 {
-	get_telemetry_mem(etmdev, EDGETPU_TELEMETRY_LOG, etmdev->log_mem);
-	get_telemetry_mem(etmdev, EDGETPU_TELEMETRY_TRACE, etmdev->trace_mem);
+	get_telemetry_mem(etmdev, GCIP_TELEMETRY_LOG, etmdev->log_mem);
+	get_telemetry_mem(etmdev, GCIP_TELEMETRY_TRACE, etmdev->trace_mem);
 }
 
 static int edgetpu_platform_setup_fw_region(struct edgetpu_mobile_platform_dev *etmdev)
@@ -259,10 +259,8 @@ static void edgetpu_platform_remove_irq(struct edgetpu_mobile_platform_dev *etmd
 
 /*
  * Fetch and set the firmware context region from device tree.
- *
- * Maybe be unused since not all chips need this.
  */
-static int __maybe_unused
+static int
 edgetpu_mobile_platform_set_fw_ctx_memory(struct edgetpu_mobile_platform_dev *etmdev)
 {
 	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
@@ -391,6 +389,11 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 
 	etdev_dbg(etdev, "Creating thermal device");
 	etdev->thermal = devm_tpu_thermal_create(etdev->dev, etdev);
+	ret = edgetpu_mobile_platform_set_fw_ctx_memory(etmdev);
+	if (ret) {
+		etdev_err(etdev, "Failed to initialize fw context memory: %d", ret);
+		goto out_destroy_fw;
+	}
 
 	if (etmdev->after_probe) {
 		ret = etmdev->after_probe(etmdev);
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index 02d560f..73e37c6 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -36,11 +36,8 @@ struct edgetpu_mobile_platform_pwr {
 	/* Block shutdown callback, may be NULL */
 	void (*block_down)(struct edgetpu_dev *etdev);
 
-	/* Chip-specific setup after the PM interface is created */
-	int (*after_create)(struct edgetpu_dev *etdev);
-
-	/* Chip-specific cleanup before the PM interface is destroyed */
-	int (*before_destroy)(struct edgetpu_dev *etdev);
+	/* After firmware is started on power up */
+	void (*post_fw_start)(struct edgetpu_dev *etdev);
 };
 
 struct edgetpu_mobile_platform_dev {
diff --git a/drivers/edgetpu/edgetpu-soc.h b/drivers/edgetpu/edgetpu-soc.h
index 00e91ef..bba617a 100644
--- a/drivers/edgetpu/edgetpu-soc.h
+++ b/drivers/edgetpu/edgetpu-soc.h
@@ -22,7 +22,11 @@ int edgetpu_soc_init(struct edgetpu_dev *etdev);
 /* Prep for running firmware: set access control, etc. */
 int edgetpu_soc_prepare_firmware(struct edgetpu_dev *etdev);
 
-/* Power management get TPU clock rate */
+/*
+ * Power management get TPU clock rate.
+ * @flags can be used by platform-specific code to pass additional flags to the SoC
+ *        handler; for calls from generic code this value must be zero.
+ */
 long edgetpu_soc_pm_get_rate(int flags);
 
 /* Power management set TPU clock rate */
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index ce4739f..9376fba 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -5,15 +5,9 @@
  * Copyright (C) 2019-2020 Google, Inc.
  */
 
-#include <linux/delay.h>
-#include <linux/dma-mapping.h>
-#include <linux/errno.h>
-#include <linux/kernel.h>
-#include <linux/mm.h>
-#include <linux/mutex.h>
-#include <linux/slab.h>
-#include <linux/spinlock.h>
-#include <linux/workqueue.h>
+#include <linux/mm_types.h>
+
+#include <gcip/gcip-telemetry.h>
 
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -21,17 +15,13 @@
 #include "edgetpu-telemetry.h"
 #include "edgetpu.h"
 
-/* When log data arrives, recheck for more log data after this delay. */
-#define TELEMETRY_LOG_RECHECK_DELAY	200	/* ms */
-
-static struct edgetpu_telemetry *
-select_telemetry(struct edgetpu_telemetry_ctx *ctx,
-		 enum edgetpu_telemetry_type type)
+static struct gcip_telemetry *select_telemetry(struct edgetpu_telemetry_ctx *ctx,
+					       enum gcip_telemetry_type type)
 {
 	switch (type) {
-	case EDGETPU_TELEMETRY_TRACE:
+	case GCIP_TELEMETRY_TRACE:
 		return &ctx->trace;
-	case EDGETPU_TELEMETRY_LOG:
+	case GCIP_TELEMETRY_LOG:
 		return &ctx->log;
 	default:
 		WARN_ONCE(1, "Unrecognized EdgeTPU telemetry type: %d", type);
@@ -40,276 +30,19 @@ select_telemetry(struct edgetpu_telemetry_ctx *ctx,
 	}
 }
 
-static int telemetry_kci(struct edgetpu_dev *etdev, struct edgetpu_telemetry *tel,
-			 int (*send_kci)(struct edgetpu_kci *, u64, u32))
-{
-	int err;
-
-	etdev_dbg(etdev, "Sending KCI %s", tel->name);
-	err = send_kci(etdev->etkci, tel->coherent_mem.tpu_addr, tel->coherent_mem.size);
-
-	if (err < 0) {
-		etdev_err(etdev, "KCI %s failed :( - %d", tel->name, err);
-		return err;
-	}
-
-	if (err > 0) {
-		etdev_err(etdev, "KCI %s returned %d", tel->name, err);
-		return -EBADMSG;
-	}
-	etdev_dbg(etdev, "KCI %s Succeeded :)", tel->name);
-	return 0;
-}
-
-static int telemetry_set_event(struct edgetpu_dev *etdev,
-			       struct edgetpu_telemetry *tel, u32 eventfd)
-{
-	struct eventfd_ctx *ctx;
-	ulong flags;
-
-	ctx = eventfd_ctx_fdget(eventfd);
-	if (IS_ERR(ctx))
-		return PTR_ERR(ctx);
-
-	write_lock_irqsave(&tel->ctx_lock, flags);
-	if (tel->ctx)
-		eventfd_ctx_put(tel->ctx);
-	tel->ctx = ctx;
-	write_unlock_irqrestore(&tel->ctx_lock, flags);
-
-	return 0;
-}
-
-static void telemetry_unset_event(struct edgetpu_dev *etdev,
-				  struct edgetpu_telemetry *tel)
-{
-	ulong flags;
-
-	write_lock_irqsave(&tel->ctx_lock, flags);
-	if (tel->ctx)
-		eventfd_ctx_put(tel->ctx);
-	tel->ctx = NULL;
-	write_unlock_irqrestore(&tel->ctx_lock, flags);
-
-	return;
-}
-
-/* Copy data out of the log buffer with wrapping */
-static void copy_with_wrap(struct edgetpu_telemetry_header *header, void *dest,
-			   u32 length, u32 size, void *start)
-{
-	const u32 wrap_bit = size + sizeof(*header);
-	u32 remaining = 0;
-	u32 head = header->head & (wrap_bit - 1);
-
-	if (head + length < size) {
-		memcpy(dest, start + head, length);
-		header->head += length;
-	} else {
-		remaining = size - head;
-		memcpy(dest, start + head, remaining);
-		memcpy(dest + remaining, start, length - remaining);
-		header->head = (header->head & wrap_bit) ^ wrap_bit;
-		header->head |= length - remaining;
-	}
-}
-
-/* Log messages from TPU CPU to dmesg */
-static void edgetpu_fw_log(struct edgetpu_telemetry *log)
-{
-	struct edgetpu_dev *etdev = log->etdev;
-	struct edgetpu_telemetry_header *header = log->header;
-	struct edgetpu_log_entry_header entry;
-	u8 *start;
-	const size_t queue_size = log->coherent_mem.size - sizeof(*header);
-	const size_t max_length = queue_size - sizeof(entry);
-	char *buffer = kmalloc(max_length + 1, GFP_ATOMIC);
-
-	if (!buffer) {
-		header->head = header->tail;
-		etdev_err_ratelimited(etdev, "failed to allocate log buffer");
-		return;
-	}
-	start = (u8 *)header + sizeof(*header);
-
-	while (header->head != header->tail) {
-		copy_with_wrap(header, &entry, sizeof(entry), queue_size,
-			       start);
-		if (entry.length == 0 || entry.length > max_length) {
-			header->head = header->tail;
-			etdev_err_ratelimited(etdev, "log queue is corrupted");
-			break;
-		}
-		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 EDGETPU_FW_LOG_LEVEL_VERBOSE:
-		case EDGETPU_FW_LOG_LEVEL_DEBUG:
-			etdev_dbg_ratelimited(etdev, "%s", buffer);
-			break;
-		case EDGETPU_FW_LOG_LEVEL_WARN:
-			etdev_warn_ratelimited(etdev, "%s", buffer);
-			break;
-		case EDGETPU_FW_LOG_LEVEL_ERROR:
-			etdev_err_ratelimited(etdev, "%s", buffer);
-			break;
-		case EDGETPU_FW_LOG_LEVEL_INFO:
-		default:
-			etdev_info_ratelimited(etdev, "%s", buffer);
-			break;
-		}
-	}
-	kfree(buffer);
-}
-
-/* Consumes the queue buffer. */
-static void edgetpu_fw_trace(struct edgetpu_telemetry *trace)
-{
-	struct edgetpu_telemetry_header *header = trace->header;
-
-	header->head = header->tail;
-}
-
-/* Worker for processing log/trace buffers. */
-
-static void telemetry_worker(struct work_struct *work)
-{
-	struct edgetpu_telemetry *tel =
-		container_of(work, struct edgetpu_telemetry, work);
-	u32 prev_head;
-	ulong flags;
-
-	/*
-	 * Loop while telemetry enabled, there is data to be consumed,
-	 * and the previous iteration made progress.  If another IRQ arrives
-	 * just after the last head != tail check we should get another worker
-	 * schedule.
-	 */
-	do {
-		spin_lock_irqsave(&tel->state_lock, flags);
-		if (tel->state != EDGETPU_TELEMETRY_ENABLED) {
-			spin_unlock_irqrestore(&tel->state_lock, flags);
-			return;
-		}
-
-		prev_head = tel->header->head;
-		if (tel->header->head != tel->header->tail) {
-			read_lock(&tel->ctx_lock);
-			if (tel->ctx)
-				eventfd_signal(tel->ctx, 1);
-			else
-				tel->fallback_fn(tel);
-			read_unlock(&tel->ctx_lock);
-		}
-
-		spin_unlock_irqrestore(&tel->state_lock, flags);
-		msleep(TELEMETRY_LOG_RECHECK_DELAY);
-	} while (tel->header->head != tel->header->tail &&
-		 tel->header->head != prev_head);
-}
-
-
-/* If the buffer queue is not empty, schedules worker. */
-static void telemetry_irq_handler(struct edgetpu_dev *etdev,
-				  struct edgetpu_telemetry *tel)
-{
-	spin_lock(&tel->state_lock);
-
-	if (tel->state == EDGETPU_TELEMETRY_ENABLED &&
-	    tel->header->head != tel->header->tail) {
-		schedule_work(&tel->work);
-	}
-
-	spin_unlock(&tel->state_lock);
-}
-
-static void telemetry_mappings_show(struct edgetpu_telemetry *tel,
-				    struct seq_file *s)
-{
-	seq_printf(s, "  %#llx %lu %s %#llx %pad\n",
-		   tel->coherent_mem.tpu_addr,
-		   DIV_ROUND_UP(tel->coherent_mem.size, PAGE_SIZE), tel->name,
-		   tel->coherent_mem.host_addr, &tel->coherent_mem.dma_addr);
-}
-
-static void telemetry_inc_mmap_count(struct edgetpu_telemetry *tel, int dif)
-{
-	mutex_lock(&tel->mmap_lock);
-	tel->mmapped_count += dif;
-	mutex_unlock(&tel->mmap_lock);
-}
-
-static int telemetry_mmap_buffer(struct edgetpu_dev *etdev,
-				 struct edgetpu_telemetry *tel,
-				 struct vm_area_struct *vma)
+static struct edgetpu_coherent_mem *select_telemetry_mem(struct edgetpu_telemetry_ctx *ctx,
+							 enum gcip_telemetry_type type)
 {
-	int ret;
-
-	mutex_lock(&tel->mmap_lock);
-
-	if (!tel->mmapped_count) {
-		ret = edgetpu_iremap_mmap(etdev, vma, &tel->coherent_mem);
-
-		if (!ret) {
-			tel->coherent_mem.host_addr = vma->vm_start;
-			tel->mmapped_count = 1;
-		}
-	} else {
-		ret = -EBUSY;
-		etdev_warn(etdev, "%s is already mmapped %ld times", tel->name,
-			   tel->mmapped_count);
+	switch (type) {
+	case GCIP_TELEMETRY_TRACE:
+		return &ctx->trace_mem;
+	case GCIP_TELEMETRY_LOG:
+		return &ctx->log_mem;
+	default:
+		WARN_ONCE(1, "Unrecognized EdgeTPU telemetry type: %d", type);
+		/* return a valid object, don't crash the kernel */
+		return &ctx->log_mem;
 	}
-
-	mutex_unlock(&tel->mmap_lock);
-
-	return ret;
-}
-
-static int telemetry_init(struct edgetpu_dev *etdev, struct edgetpu_telemetry *tel,
-			  const char *name, struct edgetpu_coherent_mem *mem, const size_t size,
-			  void (*fallback)(struct edgetpu_telemetry *))
-{
-	rwlock_init(&tel->ctx_lock);
-	tel->name = name;
-	tel->etdev = etdev;
-	tel->coherent_mem = *mem;
-
-	tel->header = (struct edgetpu_telemetry_header *)mem->vaddr;
-	tel->header->head = 0;
-	tel->header->size = 0;
-	tel->header->tail = 0;
-	tel->header->entries_dropped = 0;
-
-	tel->ctx = NULL;
-
-	spin_lock_init(&tel->state_lock);
-	INIT_WORK(&tel->work, telemetry_worker);
-	tel->fallback_fn = fallback;
-	tel->state = EDGETPU_TELEMETRY_ENABLED;
-	mutex_init(&tel->mmap_lock);
-	tel->mmapped_count = 0;
-
-	return 0;
-}
-
-static void telemetry_exit(struct edgetpu_dev *etdev,
-			   struct edgetpu_telemetry *tel)
-{
-	ulong flags;
-
-	spin_lock_irqsave(&tel->state_lock, flags);
-	/* Prevent racing with the IRQ handler or worker */
-	tel->state = EDGETPU_TELEMETRY_INVALID;
-	spin_unlock_irqrestore(&tel->state_lock, flags);
-	cancel_work_sync(&tel->work);
-
-	if (tel->ctx)
-		eventfd_ctx_put(tel->ctx);
-	tel->ctx = NULL;
 }
 
 int edgetpu_telemetry_init(struct edgetpu_dev *etdev,
@@ -319,17 +52,22 @@ int edgetpu_telemetry_init(struct edgetpu_dev *etdev,
 	int ret, i;
 
 	for (i = 0; i < etdev->num_cores; i++) {
-		ret = telemetry_init(etdev, &etdev->telemetry[i].log, "telemetry_log",
-				     log_mem ? &log_mem[i] : NULL,
-				     EDGETPU_TELEMETRY_LOG_BUFFER_SIZE, edgetpu_fw_log);
+		ret = gcip_telemetry_init(etdev->dev, &etdev->telemetry[i].log, "telemetry_log",
+					  log_mem[i].vaddr, EDGETPU_TELEMETRY_LOG_BUFFER_SIZE,
+					  gcip_telemetry_fw_log);
 		if (ret)
 			break;
+
+		etdev->telemetry[i].log_mem = log_mem[i];
+
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-		ret = telemetry_init(etdev, &etdev->telemetry[i].trace, "telemetry_trace",
-				     trace_mem ? &trace_mem[i] : NULL,
-				     EDGETPU_TELEMETRY_TRACE_BUFFER_SIZE, edgetpu_fw_trace);
+		ret = gcip_telemetry_init(etdev->dev, &etdev->telemetry[i].trace, "telemetry_trace",
+					  trace_mem[i].vaddr, EDGETPU_TELEMETRY_TRACE_BUFFER_SIZE,
+					  gcip_telemetry_fw_trace);
 		if (ret)
 			break;
+
+		etdev->telemetry[i].trace_mem = trace_mem[i];
 #endif
 	}
 
@@ -345,23 +83,34 @@ void edgetpu_telemetry_exit(struct edgetpu_dev *etdev)
 
 	for (i = 0; i < etdev->num_cores; i++) {
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-		telemetry_exit(etdev, &etdev->telemetry[i].trace);
+		gcip_telemetry_exit(&etdev->telemetry[i].trace);
 #endif
-		telemetry_exit(etdev, &etdev->telemetry[i].log);
+		gcip_telemetry_exit(&etdev->telemetry[i].log);
 	}
 }
 
 int edgetpu_telemetry_kci(struct edgetpu_dev *etdev)
 {
+	struct gcip_telemetry_kci_args log_args = {
+		.kci = etdev->etkci->kci,
+		.addr = etdev->telemetry[0].log_mem.tpu_addr,
+		.size = etdev->telemetry[0].log_mem.size,
+	};
+	struct gcip_telemetry_kci_args trace_args = {
+		.kci = etdev->etkci->kci,
+		.addr = etdev->telemetry[0].trace_mem.tpu_addr,
+		.size = etdev->telemetry[0].trace_mem.size,
+	};
 	int ret;
 
 	/* Core 0 will notify other cores. */
-	ret = telemetry_kci(etdev, &etdev->telemetry[0].log, edgetpu_kci_map_log_buffer);
+	ret = gcip_telemetry_kci(&etdev->telemetry[0].log, edgetpu_kci_map_log_buffer, &log_args);
 	if (ret)
 		return ret;
 
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-	ret = telemetry_kci(etdev, &etdev->telemetry[0].trace, edgetpu_kci_map_trace_buffer);
+	ret = gcip_telemetry_kci(&etdev->telemetry[0].trace, edgetpu_kci_map_trace_buffer,
+				 &trace_args);
 	if (ret)
 		return ret;
 #endif
@@ -369,14 +118,14 @@ int edgetpu_telemetry_kci(struct edgetpu_dev *etdev)
 	return 0;
 }
 
-int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev,
-				enum edgetpu_telemetry_type type, u32 eventfd)
+int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
+				u32 eventfd)
 {
 	int i, ret;
 
 	for (i = 0; i < etdev->num_cores; i++) {
-		ret = telemetry_set_event(etdev, select_telemetry(&etdev->telemetry[i], type),
-					  eventfd);
+		ret = gcip_telemetry_set_event(select_telemetry(&etdev->telemetry[i], type),
+					       eventfd);
 		if (ret) {
 			edgetpu_telemetry_unset_event(etdev, type);
 			return ret;
@@ -386,13 +135,12 @@ int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev,
 	return 0;
 }
 
-void edgetpu_telemetry_unset_event(struct edgetpu_dev *etdev,
-				   enum edgetpu_telemetry_type type)
+void edgetpu_telemetry_unset_event(struct edgetpu_dev *etdev, enum gcip_telemetry_type type)
 {
 	int i;
 
 	for (i = 0; i < etdev->num_cores; i++)
-		telemetry_unset_event(etdev, select_telemetry(&etdev->telemetry[i], type));
+		gcip_telemetry_unset_event(select_telemetry(&etdev->telemetry[i], type));
 }
 
 void edgetpu_telemetry_irq_handler(struct edgetpu_dev *etdev)
@@ -400,41 +148,74 @@ void edgetpu_telemetry_irq_handler(struct edgetpu_dev *etdev)
 	int i;
 
 	for (i = 0; i < etdev->num_cores; i++) {
-		telemetry_irq_handler(etdev, &etdev->telemetry[i].log);
+		gcip_telemetry_irq_handler(&etdev->telemetry[i].log);
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-		telemetry_irq_handler(etdev, &etdev->telemetry[i].trace);
+		gcip_telemetry_irq_handler(&etdev->telemetry[i].trace);
 #endif
 	}
 }
 
+static void telemetry_mappings_show(struct gcip_telemetry *tel, struct edgetpu_coherent_mem *mem,
+				    struct seq_file *s)
+{
+	seq_printf(s, "  %#llx %lu %s %#llx %pad\n", mem->tpu_addr,
+		   DIV_ROUND_UP(mem->size, PAGE_SIZE), tel->name, mem->host_addr, &mem->dma_addr);
+}
+
 void edgetpu_telemetry_mappings_show(struct edgetpu_dev *etdev,
 				     struct seq_file *s)
 {
 	int i;
 
 	for (i = 0; i < etdev->num_cores; i++) {
-		telemetry_mappings_show(&etdev->telemetry[i].log, s);
+		telemetry_mappings_show(&etdev->telemetry[i].log, &etdev->telemetry[i].log_mem, s);
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-		telemetry_mappings_show(&etdev->telemetry[i].trace, s);
+		telemetry_mappings_show(&etdev->telemetry[i].trace, &etdev->telemetry[i].trace_mem,
+					s);
 #endif
 	}
 }
 
-int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+struct edgetpu_telemetry_mmap_args {
+	struct edgetpu_dev *etdev;
+	struct edgetpu_coherent_mem *mem;
+	struct vm_area_struct *vma;
+};
+
+static int telemetry_mmap_buffer(void *args)
+{
+	struct edgetpu_telemetry_mmap_args *data = args;
+	int ret;
+
+	ret = edgetpu_iremap_mmap(data->etdev, data->vma, data->mem);
+
+	if (!ret)
+		data->mem->host_addr = data->vma->vm_start;
+
+	return ret;
+}
+
+int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				  struct vm_area_struct *vma, int core_id)
 {
-	return telemetry_mmap_buffer(etdev, select_telemetry(&etdev->telemetry[core_id], type),
-				     vma);
+	struct edgetpu_telemetry_mmap_args args = {
+		.etdev = etdev,
+		.mem = select_telemetry_mem(&etdev->telemetry[core_id], type),
+		.vma = vma,
+	};
+
+	return gcip_telemetry_mmap_buffer(select_telemetry(&etdev->telemetry[core_id], type),
+					  telemetry_mmap_buffer, &args);
 }
 
-void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				      int core_id)
 {
-	telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), 1);
+	gcip_telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), 1);
 }
 
-void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				      int core_id)
 {
-	telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), -1);
+	gcip_telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), -1);
 }
diff --git a/drivers/edgetpu/edgetpu-telemetry.h b/drivers/edgetpu/edgetpu-telemetry.h
index b703299..26883bd 100644
--- a/drivers/edgetpu/edgetpu-telemetry.h
+++ b/drivers/edgetpu/edgetpu-telemetry.h
@@ -7,87 +7,23 @@
 #ifndef __EDGETPU_TELEMETRY_H__
 #define __EDGETPU_TELEMETRY_H__
 
-#include <linux/eventfd.h>
-#include <linux/mm.h>
-#include <linux/mutex.h>
+#include <linux/mm_types.h>
 #include <linux/seq_file.h>
-#include <linux/spinlock.h>
-#include <linux/types.h>
-#include <linux/workqueue.h>
+
+#include <gcip/gcip-telemetry.h>
 
 #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)
-
 /* Buffer size must be a power of 2 */
 #define EDGETPU_TELEMETRY_LOG_BUFFER_SIZE (16 * 4096)
 #define EDGETPU_TELEMETRY_TRACE_BUFFER_SIZE (64 * 4096)
 
-enum edgetpu_telemetry_state {
-	EDGETPU_TELEMETRY_DISABLED = 0,
-	EDGETPU_TELEMETRY_ENABLED = 1,
-	EDGETPU_TELEMETRY_INVALID = -1,
-};
-
-/* To specify the target of operation. */
-enum edgetpu_telemetry_type {
-	EDGETPU_TELEMETRY_LOG = 0,
-	EDGETPU_TELEMETRY_TRACE = 1,
-};
-
-struct edgetpu_telemetry_header {
-	u32 head;
-	u32 size;
-	u32 reserved0[14]; /* Place head and tail into different cache lines */
-	u32 tail;
-	u32 entries_dropped; /* Number of entries dropped due to buffer full */
-	u32 reserved1[14]; /* Pad to 128 bytes in total */
-};
-
-struct edgetpu_log_entry_header {
-	s16 code;
-	u16 length;
-	u64 timestamp;
-	u16 crc16;
-} __packed;
-
-struct edgetpu_telemetry {
-	struct edgetpu_dev *etdev;
-
-	/*
-	 * State transitioning is to prevent racing in IRQ handlers. e.g. the
-	 * interrupt comes when the kernel is releasing buffers.
-	 */
-	enum edgetpu_telemetry_state state;
-	spinlock_t state_lock; /* protects state */
-
-	struct edgetpu_coherent_mem coherent_mem;
-	struct edgetpu_telemetry_header *header;
-
-	struct eventfd_ctx *ctx; /* signal this to notify the runtime */
-	rwlock_t ctx_lock; /* protects ctx */
-	const char *name; /* for debugging */
-
-	/* Worker for handling data. */
-	struct work_struct work;
-	/* Fallback function to call for default log/trace handling. */
-	void (*fallback_fn)(struct edgetpu_telemetry *tel);
-	struct mutex mmap_lock; /* protects mmapped_count */
-	/* number of VMAs that are mapped to this telemetry buffer */
-	long mmapped_count;
-};
-
 struct edgetpu_telemetry_ctx {
-	struct edgetpu_telemetry log;
-	struct edgetpu_telemetry trace;
+	struct gcip_telemetry log;
+	struct edgetpu_coherent_mem log_mem;
+	struct gcip_telemetry trace;
+	struct edgetpu_coherent_mem trace_mem;
 };
 
 /*
@@ -119,11 +55,10 @@ int edgetpu_telemetry_kci(struct edgetpu_dev *etdev);
  *
  * Returns 0 on success, or a negative errno on error.
  */
-int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev,
-				enum edgetpu_telemetry_type type, u32 eventfd);
+int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
+				u32 eventfd);
 /* Removes previously set event. */
-void edgetpu_telemetry_unset_event(struct edgetpu_dev *etdev,
-				   enum edgetpu_telemetry_type type);
+void edgetpu_telemetry_unset_event(struct edgetpu_dev *etdev, enum gcip_telemetry_type type);
 
 /* Checks telemetries and signals eventfd if needed. */
 void edgetpu_telemetry_irq_handler(struct edgetpu_dev *etdev);
@@ -133,11 +68,11 @@ void edgetpu_telemetry_mappings_show(struct edgetpu_dev *etdev,
 				     struct seq_file *s);
 
 /* Map telemetry buffer into user space. */
-int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				  struct vm_area_struct *vma, int core_id);
-void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				      int core_id);
-void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
+void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev, enum gcip_telemetry_type type,
 				      int core_id);
 
 #endif /* __EDGETPU_TELEMETRY_H__ */
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index 1b988e4..7f6d2f0 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -6,7 +6,8 @@
 CONFIG_GCIP		?= m
 obj-$(CONFIG_GCIP)	+= gcip.o
 
-gcip-objs		:= gcip-firmware.o gcip-image-config.o gcip-kci.o gcip-mailbox.o gcip-mem-pool.o
+gcip-objs		:= gcip-firmware.o gcip-image-config.o gcip-kci.o gcip-mailbox.o \
+			   gcip-mem-pool.o gcip-telemetry.o
 
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-telemetry.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-telemetry.c
new file mode 100644
index 0000000..af11169
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-telemetry.c
@@ -0,0 +1,260 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * GCIP telemetry: logging and tracing.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#include <linux/delay.h>
+#include <linux/dev_printk.h>
+#include <linux/eventfd.h>
+#include <linux/mutex.h>
+#include <linux/slab.h>
+#include <linux/spinlock.h>
+#include <linux/workqueue.h>
+
+#include <gcip/gcip-telemetry.h>
+
+int gcip_telemetry_kci(struct gcip_telemetry *tel,
+		       int (*send_kci)(struct gcip_telemetry_kci_args *),
+		       struct gcip_telemetry_kci_args *args)
+{
+	int err;
+
+	dev_dbg(tel->dev, "Sending KCI %s", tel->name);
+	err = send_kci(args);
+
+	if (err < 0) {
+		dev_err(tel->dev, "KCI %s failed - %d", tel->name, err);
+		return err;
+	}
+
+	if (err > 0) {
+		dev_err(tel->dev, "KCI %s returned %d", tel->name, err);
+		return -EBADMSG;
+	}
+
+	dev_dbg(tel->dev, "KCI %s Succeeded", tel->name);
+
+	return 0;
+}
+
+int gcip_telemetry_set_event(struct gcip_telemetry *tel, u32 eventfd)
+{
+	struct eventfd_ctx *ctx;
+	ulong flags;
+
+	ctx = eventfd_ctx_fdget(eventfd);
+	if (IS_ERR(ctx))
+		return PTR_ERR(ctx);
+
+	write_lock_irqsave(&tel->ctx_lock, flags);
+	if (tel->ctx)
+		eventfd_ctx_put(tel->ctx);
+	tel->ctx = ctx;
+	write_unlock_irqrestore(&tel->ctx_lock, flags);
+
+	return 0;
+}
+
+void gcip_telemetry_unset_event(struct gcip_telemetry *tel)
+{
+	ulong flags;
+
+	write_lock_irqsave(&tel->ctx_lock, flags);
+	if (tel->ctx)
+		eventfd_ctx_put(tel->ctx);
+	tel->ctx = NULL;
+	write_unlock_irqrestore(&tel->ctx_lock, flags);
+}
+
+/* Copy data out of the log buffer with wrapping. */
+static void copy_with_wrap(struct gcip_telemetry_header *header, void *dest, u32 length, u32 size,
+			   void *start)
+{
+	const u32 wrap_bit = size + sizeof(*header);
+	u32 remaining = 0;
+	u32 head = header->head & (wrap_bit - 1);
+
+	if (head + length < size) {
+		memcpy(dest, start + head, length);
+		header->head += length;
+	} else {
+		remaining = size - head;
+		memcpy(dest, start + head, remaining);
+		memcpy(dest + remaining, start, length - remaining);
+		header->head = (header->head & wrap_bit) ^ wrap_bit;
+		header->head |= length - remaining;
+	}
+}
+
+void gcip_telemetry_fw_log(struct gcip_telemetry *log)
+{
+	struct device *dev = log->dev;
+	struct gcip_telemetry_header *header = log->header;
+	struct gcip_log_entry_header entry;
+	u8 *start;
+	const size_t queue_size = header->size - sizeof(*header);
+	const size_t max_length = queue_size - sizeof(entry);
+	char *buffer = kmalloc(max_length + 1, GFP_ATOMIC);
+
+	if (!buffer) {
+		header->head = header->tail;
+		dev_err(dev, "failed to allocate log buffer");
+		return;
+	}
+	start = (u8 *)header + sizeof(*header);
+
+	while (header->head != header->tail) {
+		copy_with_wrap(header, &entry, sizeof(entry), queue_size, start);
+		if (entry.length == 0 || entry.length > max_length) {
+			header->head = header->tail;
+			dev_err(dev, "log queue is corrupted");
+			break;
+		}
+		copy_with_wrap(header, buffer, entry.length, queue_size, start);
+		buffer[entry.length] = 0;
+
+		if (entry.code > GCIP_FW_DMESG_LOG_LEVEL)
+			continue;
+
+		switch (entry.code) {
+		case GCIP_FW_LOG_LEVEL_VERBOSE:
+		case GCIP_FW_LOG_LEVEL_DEBUG:
+			dev_dbg(dev, "%s", buffer);
+			break;
+		case GCIP_FW_LOG_LEVEL_WARN:
+			dev_warn(dev, "%s", buffer);
+			break;
+		case GCIP_FW_LOG_LEVEL_ERROR:
+			dev_err(dev, "%s", buffer);
+			break;
+		case GCIP_FW_LOG_LEVEL_INFO:
+		default:
+			dev_info(dev, "%s", buffer);
+			break;
+		}
+	}
+	kfree(buffer);
+}
+
+void gcip_telemetry_fw_trace(struct gcip_telemetry *trace)
+{
+	struct gcip_telemetry_header *header = trace->header;
+
+	header->head = header->tail;
+}
+
+void gcip_telemetry_irq_handler(struct gcip_telemetry *tel)
+{
+	spin_lock(&tel->state_lock);
+
+	if (tel->state == GCIP_TELEMETRY_ENABLED && tel->header->head != tel->header->tail)
+		schedule_work(&tel->work);
+
+	spin_unlock(&tel->state_lock);
+}
+
+void gcip_telemetry_inc_mmap_count(struct gcip_telemetry *tel, int dif)
+{
+	mutex_lock(&tel->mmap_lock);
+	tel->mmapped_count += dif;
+	mutex_unlock(&tel->mmap_lock);
+}
+
+int gcip_telemetry_mmap_buffer(struct gcip_telemetry *tel, int (*mmap)(void *), void *args)
+{
+	int ret;
+
+	mutex_lock(&tel->mmap_lock);
+
+	if (!tel->mmapped_count) {
+		ret = mmap(args);
+
+		if (!ret)
+			tel->mmapped_count = 1;
+	} else {
+		ret = -EBUSY;
+		dev_warn(tel->dev, "%s is already mmapped %ld times", tel->name,
+			 tel->mmapped_count);
+	}
+
+	mutex_unlock(&tel->mmap_lock);
+
+	return ret;
+}
+
+/* Worker for processing log/trace buffers. */
+static void gcip_telemetry_worker(struct work_struct *work)
+{
+	struct gcip_telemetry *tel = container_of(work, struct gcip_telemetry, work);
+	u32 prev_head;
+	ulong flags;
+
+	/*
+	 * Loops while telemetry enabled, there is data to be consumed, and the previous iteration
+	 * made progress. If another IRQ arrives just after the last head != tail check we should
+	 * get another worker schedule.
+	 */
+	do {
+		spin_lock_irqsave(&tel->state_lock, flags);
+		if (tel->state != GCIP_TELEMETRY_ENABLED) {
+			spin_unlock_irqrestore(&tel->state_lock, flags);
+			return;
+		}
+
+		prev_head = tel->header->head;
+		if (tel->header->head != tel->header->tail) {
+			read_lock(&tel->ctx_lock);
+			if (tel->ctx)
+				eventfd_signal(tel->ctx, 1);
+			else
+				tel->fallback_fn(tel);
+			read_unlock(&tel->ctx_lock);
+		}
+
+		spin_unlock_irqrestore(&tel->state_lock, flags);
+		msleep(GCIP_TELEMETRY_LOG_RECHECK_DELAY);
+	} while (tel->header->head != tel->header->tail && tel->header->head != prev_head);
+}
+
+int gcip_telemetry_init(struct device *dev, struct gcip_telemetry *tel, const char *name,
+			void *vaddr, const size_t size,
+			void (*fallback_fn)(struct gcip_telemetry *))
+{
+	rwlock_init(&tel->ctx_lock);
+	tel->name = name;
+	tel->dev = dev;
+
+	tel->header = (struct gcip_telemetry_header *)vaddr;
+	tel->header->head = 0;
+	tel->header->tail = 0;
+	tel->header->size = size;
+	tel->header->entries_dropped = 0;
+
+	tel->ctx = NULL;
+
+	spin_lock_init(&tel->state_lock);
+	INIT_WORK(&tel->work, gcip_telemetry_worker);
+	tel->fallback_fn = fallback_fn;
+	tel->state = GCIP_TELEMETRY_ENABLED;
+	mutex_init(&tel->mmap_lock);
+	tel->mmapped_count = 0;
+
+	return 0;
+}
+
+void gcip_telemetry_exit(struct gcip_telemetry *tel)
+{
+	ulong flags;
+
+	spin_lock_irqsave(&tel->state_lock, flags);
+	/* Prevents racing with the IRQ handler or worker. */
+	tel->state = GCIP_TELEMETRY_INVALID;
+	spin_unlock_irqrestore(&tel->state_lock, flags);
+	cancel_work_sync(&tel->work);
+
+	if (tel->ctx)
+		eventfd_ctx_put(tel->ctx);
+	tel->ctx = NULL;
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-telemetry.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-telemetry.h
new file mode 100644
index 0000000..4556291
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-telemetry.h
@@ -0,0 +1,123 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * GCIP telemetry: logging and tracing.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#ifndef __GCIP_TELEMETRY_H__
+#define __GCIP_TELEMETRY_H__
+
+#include <linux/device.h>
+#include <linux/eventfd.h>
+#include <linux/mutex.h>
+#include <linux/rwlock_types.h>
+#include <linux/spinlock.h>
+#include <linux/types.h>
+#include <linux/workqueue.h>
+
+/* Log level codes used by gcip firmware. */
+#define GCIP_FW_LOG_LEVEL_VERBOSE (2)
+#define GCIP_FW_LOG_LEVEL_DEBUG (1)
+#define GCIP_FW_LOG_LEVEL_INFO (0)
+#define GCIP_FW_LOG_LEVEL_WARN (-1)
+#define GCIP_FW_LOG_LEVEL_ERROR (-2)
+
+#define GCIP_FW_DMESG_LOG_LEVEL (GCIP_FW_LOG_LEVEL_WARN)
+
+/* When log data arrives, recheck for more log data after this delay. */
+#define GCIP_TELEMETRY_LOG_RECHECK_DELAY 200 /* ms */
+
+enum gcip_telemetry_state {
+	GCIP_TELEMETRY_DISABLED = 0,
+	GCIP_TELEMETRY_ENABLED = 1,
+	GCIP_TELEMETRY_INVALID = -1,
+};
+
+/* To specify the target of operation. */
+enum gcip_telemetry_type {
+	GCIP_TELEMETRY_LOG = 0,
+	GCIP_TELEMETRY_TRACE = 1,
+};
+
+struct gcip_telemetry_header {
+	u32 head;
+	u32 size;
+	u32 reserved0[14]; /* Place head and tail into different cache lines */
+	u32 tail;
+	u32 entries_dropped; /* Number of entries dropped due to buffer full */
+	u32 reserved1[14]; /* Pad to 128 bytes in total */
+};
+
+struct gcip_log_entry_header {
+	s16 code;
+	u16 length;
+	u64 timestamp;
+	u16 crc16;
+} __packed;
+
+struct gcip_telemetry {
+	/* Device used for logging and memory allocation. */
+	struct device *dev;
+
+	/*
+	 * State transitioning is to prevent racing in IRQ handlers. e.g. the interrupt comes when
+	 * the kernel is releasing buffers.
+	 */
+	enum gcip_telemetry_state state;
+	spinlock_t state_lock; /* protects state */
+
+	struct gcip_telemetry_header *header;
+
+	struct eventfd_ctx *ctx; /* signal this to notify the runtime */
+	rwlock_t ctx_lock; /* protects ctx */
+	const char *name; /* for debugging */
+
+	struct work_struct work; /* worker for handling data */
+	/* Fallback function to call for default log/trace handling. */
+	void (*fallback_fn)(struct gcip_telemetry *tel);
+	struct mutex mmap_lock; /* protects mmapped_count */
+	long mmapped_count; /* number of VMAs that are mapped to this telemetry buffer */
+};
+
+struct gcip_kci;
+
+struct gcip_telemetry_kci_args {
+	struct gcip_kci *kci;
+	u64 addr;
+	u32 size;
+};
+
+/* Sends telemetry KCI through send kci callback and args. */
+int gcip_telemetry_kci(struct gcip_telemetry *tel,
+		       int (*send_kci)(struct gcip_telemetry_kci_args *),
+		       struct gcip_telemetry_kci_args *args);
+/* Sets the eventfd for telemetry. */
+int gcip_telemetry_set_event(struct gcip_telemetry *tel, u32 eventfd);
+/* Unsets the eventfd for telemetry. */
+void gcip_telemetry_unset_event(struct gcip_telemetry *tel);
+/* Fallback to log messages from host CPU to dmesg. */
+void gcip_telemetry_fw_log(struct gcip_telemetry *log);
+/* Fallback to consumes the trace buffer. */
+void gcip_telemetry_fw_trace(struct gcip_telemetry *trace);
+/* Interrupt handler to schedule the worker when the buffer is not empty. */
+void gcip_telemetry_irq_handler(struct gcip_telemetry *tel);
+/* Increases the telemetry mmap count. */
+void gcip_telemetry_inc_mmap_count(struct gcip_telemetry *tel, int dif);
+/* Mmaps the telemetry buffer through mmap callback and args. */
+int gcip_telemetry_mmap_buffer(struct gcip_telemetry *tel, int (*mmap)(void *), void *args);
+/*
+ * Initializes struct gcip_telemetry.
+ *
+ * @vaddr: Virtual address of the queue buffer.
+ * @size: Size of the queue buffer. Must be power of 2 and greater than the size of struct
+ * gcip_telemetry_header.
+ * @fallback_fn: Fallback function to call for default log/trace handling.
+ */
+int gcip_telemetry_init(struct device *dev, struct gcip_telemetry *tel, const char *name,
+			void *vaddr, const size_t size,
+			void (*fallback_fn)(struct gcip_telemetry *));
+/* Exits and sets the telemetry state to GCIP_TELEMETRY_INVALID. */
+void gcip_telemetry_exit(struct gcip_telemetry *tel);
+
+#endif /* __GCIP_TELEMETRY_H__ */
diff --git a/drivers/edgetpu/mobile-firmware.c b/drivers/edgetpu/mobile-firmware.c
index f9ae004..deeb15c 100644
--- a/drivers/edgetpu/mobile-firmware.c
+++ b/drivers/edgetpu/mobile-firmware.c
@@ -313,52 +313,6 @@ out:
 	return ret;
 }
 
-/* Load firmware for chips that use carveout memory for a single chip. */
-int edgetpu_firmware_chip_load_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc, const char *name)
-{
-	int ret;
-	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct device *dev = etdev->dev;
-	const struct firmware *fw;
-	size_t aligned_size;
-
-	ret = request_firmware(&fw, name, dev);
-	if (ret) {
-		etdev_dbg(etdev,
-			  "%s: request '%s' failed: %d\n", __func__, name, ret);
-		return ret;
-	}
-
-	aligned_size = ALIGN(fw->size, fw_desc->buf.used_size_align);
-	if (aligned_size > fw_desc->buf.alloc_size) {
-		etdev_dbg(etdev,
-			   "%s: firmware buffer too small: alloc size=%#zx, required size=%#zx\n",
-			  __func__, fw_desc->buf.alloc_size, aligned_size);
-		ret = -ENOSPC;
-		goto out_release_firmware;
-	}
-
-	memcpy(fw_desc->buf.vaddr, fw->data, fw->size);
-	fw_desc->buf.used_size = aligned_size;
-	/* May return NULL on out of memory, driver must handle properly */
-	fw_desc->buf.name = kstrdup(name, GFP_KERNEL);
-
-out_release_firmware:
-	release_firmware(fw);
-	return ret;
-}
-
-void edgetpu_firmware_chip_unload_locked(
-		struct edgetpu_firmware *et_fw,
-		struct edgetpu_firmware_desc *fw_desc)
-{
-	kfree(fw_desc->buf.name);
-	fw_desc->buf.name = NULL;
-	fw_desc->buf.used_size = 0;
-}
-
 int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset)
 {
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 4d19b2a..04caa01 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -10,7 +10,6 @@
 #include <linux/gsa/gsa_tpu.h>
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
-#include <soc/google/bcl.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-firmware.h"
@@ -21,7 +20,6 @@
 #include "edgetpu-pm.h"
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
-#include "mobile-soc-gsx01.h"
 
 #include "edgetpu-pm.c"
 #include "edgetpu-soc.h"
@@ -301,14 +299,9 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 	if (ret) {
 		mobile_power_down(etpm);
 	} else {
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-		if (!etdev->soc_data->bcl_dev)
-			etdev->soc_data->bcl_dev = google_retrieve_bcl_handle();
-		if (etdev->soc_data->bcl_dev)
-			google_init_tpu_ratio(etdev->soc_data->bcl_dev);
-#endif
+		if (platform_pwr->post_fw_start)
+			platform_pwr->post_fw_start(etdev);
 	}
-
 	return ret;
 }
 
@@ -408,18 +401,15 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 	if (IS_ERR_OR_NULL(platform_pwr->debugfs_dir)) {
 		dev_warn(etdev->dev, "Failed to create debug FS power");
 		/* don't fail the procedure on debug FS creation fails */
-		return 0;
+	} else {
+		debugfs_create_file("state", 0660, platform_pwr->debugfs_dir, etdev,
+				    &fops_tpu_pwr_state);
+		debugfs_create_file("min_state", 0660, platform_pwr->debugfs_dir, etdev,
+				    &fops_tpu_min_pwr_state);
+		debugfs_create_file("policy", 0660, platform_pwr->debugfs_dir, etdev,
+				    &fops_tpu_pwr_policy);
 	}
-	debugfs_create_file("state", 0660, platform_pwr->debugfs_dir, etdev, &fops_tpu_pwr_state);
-	debugfs_create_file("min_state", 0660, platform_pwr->debugfs_dir, etdev,
-			    &fops_tpu_min_pwr_state);
-	debugfs_create_file("policy", 0660, platform_pwr->debugfs_dir, etdev, &fops_tpu_pwr_policy);
-	edgetpu_soc_pm_init(etdev);
-
-	if (platform_pwr->after_create)
-		ret = platform_pwr->after_create(etdev);
-
-	return ret;
+	return edgetpu_soc_pm_init(etdev);
 }
 
 static void mobile_pm_before_destroy(struct edgetpu_pm *etpm)
@@ -428,9 +418,6 @@ static void mobile_pm_before_destroy(struct edgetpu_pm *etpm)
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 
-	if (platform_pwr->before_destroy)
-		platform_pwr->before_destroy(etdev);
-
 	debugfs_remove_recursive(platform_pwr->debugfs_dir);
 	pm_runtime_disable(etpm->etdev->dev);
 	edgetpu_soc_pm_exit(etdev);
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index d69cd32..d8a2535 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -5,7 +5,6 @@
  * Copyright (C) 2022 Google LLC
  */
 
-#include <linux/acpm_dvfs.h>
 #include <linux/device.h>
 #include <linux/gsa/gsa_tpu.h>
 #include <linux/platform_device.h>
@@ -15,6 +14,9 @@
 #include <soc/google/exynos_pm_qos.h>
 #include <soc/google/gs_tmu.h>
 
+/* TODO(b/240363978): Use system ACPM header once it's ready. */
+#include "include/linux/acpm_dvfs.h"
+
 #include "edgetpu-internal.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-kci.h"
diff --git a/drivers/edgetpu/mobile-soc-gsx01.h b/drivers/edgetpu/mobile-soc-gsx01.h
index c402d8d..d07d1d4 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.h
+++ b/drivers/edgetpu/mobile-soc-gsx01.h
@@ -11,11 +11,8 @@
 #include <linux/kernel.h>
 #include <linux/mutex.h>
 #include <linux/types.h>
-#include <soc/google/exynos_pm_qos.h>
-
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
 #include <soc/google/bcl.h>
-#endif
+#include <soc/google/exynos_pm_qos.h>
 
 /* SoC data for GSx01 platforms */
 struct edgetpu_soc_data {
@@ -30,9 +27,7 @@ struct edgetpu_soc_data {
 	unsigned int performance_scenario;
 	int scenario_count;
 	struct mutex scenario_lock;
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
 	struct bcl_device *bcl_dev;
-#endif
 };
 
 /*
diff --git a/drivers/edgetpu/rio-platform.c b/drivers/edgetpu/rio-platform.c
index f186bff..7c65d5d 100644
--- a/drivers/edgetpu/rio-platform.c
+++ b/drivers/edgetpu/rio-platform.c
@@ -104,12 +104,6 @@ static int rio_parse_set_dt_property(struct edgetpu_mobile_platform_dev *etmdev)
 	struct device *dev = etdev->dev;
 	int ret;
 
-	ret = edgetpu_mobile_platform_set_fw_ctx_memory(etmdev);
-	if (ret) {
-		etdev_err(etdev, "Failed to initialize fw context memory: %d", ret);
-		return ret;
-	}
-
 	ret = rio_mmu_set_shareability(dev);
 	if (ret)
 		etdev_warn(etdev, "failed to enable shareability: %d", ret);
diff --git a/drivers/edgetpu/rio-pm.c b/drivers/edgetpu/rio-pm.c
index 9696843..29e5dbf 100644
--- a/drivers/edgetpu/rio-pm.c
+++ b/drivers/edgetpu/rio-pm.c
@@ -7,11 +7,13 @@
 
 #include <linux/delay.h>
 #include <linux/iopoll.h>
+#include <soc/google/bcl.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-mobile-platform.h"
 #include "edgetpu-soc.h"
+#include "mobile-soc-gsx01.h"
 #include "mobile-pm.h"
 
 #define TPU_DEFAULT_POWER_STATE TPU_ACTIVE_NOM
@@ -159,6 +161,14 @@ static void rio_block_down(struct edgetpu_dev *etdev)
 		etdev_warn(etdev, "blk_shutdown timeout\n");
 }
 
+static void rio_post_fw_start(struct edgetpu_dev *etdev)
+{
+	if (!etdev->soc_data->bcl_dev)
+		etdev->soc_data->bcl_dev = google_retrieve_bcl_handle();
+	if (etdev->soc_data->bcl_dev)
+		google_init_tpu_ratio(etdev->soc_data->bcl_dev);
+}
+
 int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
@@ -167,6 +177,7 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 	platform_pwr->lpm_up = rio_lpm_up;
 	platform_pwr->lpm_down = rio_lpm_down;
 	platform_pwr->block_down = rio_block_down;
+	platform_pwr->post_fw_start = rio_post_fw_start;
 
 	return edgetpu_mobile_pm_create(etdev);
 }