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

* whitechapel:
  edgetpu: protect tel is_mapped with a mutex lock
  edgetpu: check size more strictly on mbox attr validation
  edgetpu: save IRQ state for KCI spinlocks
  edgetpu: add NOWARN flag when pinning user pages
  edgetpu: abrolhos: only throttle to active TPU states
  edgetpu: add reverse KCI handler for firmware crash events
  edgetpu: KCI update response data structure comments for Reverse KCI
  edgetpu: make edgetpu_usage_get_counter static
  edgetpu: remove deprecated FW info KCI
  edgetpu: unset all per die events on client remove
  edgetpu: allow one process to mmap log/trace buffers
  edgetpu: add "counters" usage stats
  edgetpu: Alloc dump memory outside iremap region
  edgetpu: abrolhos: add min_state debugfs entry

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

19 files changed, 513 insertions, 180 deletions

diff --git a/drivers/edgetpu/abrolhos-debug-dump.c b/drivers/edgetpu/abrolhos-debug-dump.c
index cdc57e2..a4dd732 100644
--- a/drivers/edgetpu/abrolhos-debug-dump.c
+++ b/drivers/edgetpu/abrolhos-debug-dump.c
@@ -92,11 +92,10 @@ int edgetpu_debug_dump_init(struct edgetpu_dev *etdev)
 	size = EDGETPU_DEBUG_DUMP_MEM_SIZE;
 
 	/*
-	 * Allocate buffers for various dump segments and map them to FW
-	 * accessible regions
+	 * Allocate a buffer for various dump segments
 	 */
-	ret = edgetpu_iremap_alloc(etdev, size, &etdev->debug_dump_mem,
-				   EDGETPU_CONTEXT_KCI);
+	ret = edgetpu_alloc_coherent(etdev, size, &etdev->debug_dump_mem,
+				     EDGETPU_CONTEXT_KCI);
 	if (ret) {
 		etdev_err(etdev, "Debug dump seg alloc failed");
 		etdev->debug_dump_mem.vaddr = NULL;
@@ -130,7 +129,7 @@ void edgetpu_debug_dump_exit(struct edgetpu_dev *etdev)
 	/*
 	 * Free the memory assigned for debug dump
 	 */
-	edgetpu_iremap_free(etdev, &etdev->debug_dump_mem,
-			    EDGETPU_CONTEXT_KCI);
+	edgetpu_free_coherent(etdev, &etdev->debug_dump_mem,
+			      EDGETPU_CONTEXT_KCI);
 	kfree(etdev->debug_dump_handlers);
 }
diff --git a/drivers/edgetpu/abrolhos-platform.h b/drivers/edgetpu/abrolhos-platform.h
index ecd3742..721f24d 100644
--- a/drivers/edgetpu/abrolhos-platform.h
+++ b/drivers/edgetpu/abrolhos-platform.h
@@ -22,6 +22,9 @@
 struct edgetpu_platform_pwr {
 	struct mutex policy_lock;
 	enum tpu_pwr_state curr_policy;
+	struct mutex state_lock;
+	u64 min_state;
+	u64 requested_state;
 };
 
 struct abrolhos_platform_dev {
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 0b0903b..07f0369 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -24,7 +24,6 @@
 
 #include "soc/google/exynos_pm_qos.h"
 #include "soc/google/bts.h"
-#include "soc/google/bcl.h"
 
 #include "edgetpu-pm.c"
 
@@ -82,11 +81,12 @@ static int abrolhos_pwr_state_init(struct device *dev)
 	return ret;
 }
 
-static int abrolhos_pwr_state_set(void *data, u64 val)
+static int abrolhos_pwr_state_set_locked(void *data, u64 val)
 {
 	int ret;
 	int curr_state;
-	struct device *dev = (struct device *)data;
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct device *dev = etdev->dev;
 
 	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 
@@ -132,9 +132,10 @@ static int abrolhos_pwr_state_set(void *data, u64 val)
 	return ret;
 }
 
-static int abrolhos_pwr_state_get(void *data, u64 *val)
+static int abrolhos_pwr_state_get_locked(void *data, u64 *val)
 {
-	struct device *dev = (struct device *)data;
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct device *dev = etdev->dev;
 
 	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 	dev_dbg(dev, "current tpu state: %llu\n", *val);
@@ -142,6 +143,61 @@ static int abrolhos_pwr_state_get(void *data, u64 *val)
 	return 0;
 }
 
+static int abrolhos_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->requested_state = val;
+	if (val >= platform_pwr->min_state)
+		ret = abrolhos_pwr_state_set_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int abrolhos_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret;
+
+	mutex_lock(&platform_pwr->state_lock);
+	ret = abrolhos_pwr_state_get_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int abrolhos_min_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->min_state = val;
+	if (val >= platform_pwr->requested_state)
+		ret = abrolhos_pwr_state_set_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int abrolhos_min_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+
+	mutex_lock(&platform_pwr->state_lock);
+	*val = platform_pwr->min_state;
+	mutex_unlock(&platform_pwr->state_lock);
+	return 0;
+}
+
 static int abrolhos_pwr_policy_set(void *data, u64 val)
 {
 	struct abrolhos_platform_dev *edgetpu_pdev = (typeof(edgetpu_pdev))data;
@@ -181,6 +237,9 @@ DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_policy, abrolhos_pwr_policy_get,
 DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, abrolhos_pwr_state_get,
 			 abrolhos_pwr_state_set, "%llu\n");
 
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, abrolhos_min_pwr_state_get,
+			 abrolhos_min_pwr_state_set, "%llu\n");
+
 static int edgetpu_core_rate_get(void *data, u64 *val)
 {
 	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN,
@@ -401,9 +460,8 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 {
 	struct edgetpu_dev *etdev = etpm->etdev;
 	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
-	struct device *dev = etdev->dev;
-	int ret = abrolhos_pwr_state_set(dev,
-					 abrolhos_get_initial_pwr_state(dev));
+	int ret = abrolhos_pwr_state_set(
+		etpm->etdev, abrolhos_get_initial_pwr_state(etdev->dev));
 	enum edgetpu_firmware_status firmware_status;
 
 	etdev_info(etpm->etdev, "Powering up\n");
@@ -466,12 +524,6 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 
 	if (ret)
 		abrolhos_power_down(etpm);
-	else {
-        	if (!etdev->bcl_dev)
-			etdev->bcl_dev = gs101_retrieve_bcl_handle();
-		if (etdev->bcl_dev)
-			gs101_init_tpu_ratio(etdev->bcl_dev);
-	}
 
 	return ret;
 }
@@ -532,16 +584,23 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm)
 	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
 	u64 val;
 	int res;
+	int min_state = edgetpu_pdev->platform_pwr.min_state;
 
 	etdev_info(etdev, "Powering down\n");
 
+	if (min_state >= TPU_DEEP_SLEEP_CLOCKS_SLOW) {
+		etdev_info(etdev, "Power down skipped due to min state = %d\n",
+			   min_state);
+		return;
+	}
+
 	/* Remove our vote for INT/MIF state (if any) */
 	exynos_pm_qos_update_request(&int_min, 0);
 	exynos_pm_qos_update_request(&mif_min, 0);
 
 	abrolhos_pm_cleanup_bts_scenario(etdev);
 
-	if (abrolhos_pwr_state_get(etdev->dev, &val)) {
+	if (abrolhos_pwr_state_get(etdev, &val)) {
 		etdev_warn(etdev, "Failed to read current power state\n");
 		val = TPU_ACTIVE_NOM;
 	}
@@ -561,7 +620,7 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm)
 	res = gsa_send_tpu_cmd(edgetpu_pdev->gsa_dev, GSA_TPU_SHUTDOWN);
 	if (res < 0)
 		etdev_warn(etdev, "GSA shutdown request failed (%d)\n", res);
-	abrolhos_pwr_state_set(etdev->dev, TPU_OFF);
+	abrolhos_pwr_state_set(etdev, TPU_OFF);
 }
 
 static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
@@ -575,11 +634,13 @@ static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
 	if (ret)
 		return ret;
 
-	ret = abrolhos_pwr_state_set(dev, abrolhos_get_initial_pwr_state(dev));
+	ret = abrolhos_pwr_state_set(etdev,
+				     abrolhos_get_initial_pwr_state(dev));
 	if (ret)
 		return ret;
 
 	mutex_init(&edgetpu_pdev->platform_pwr.policy_lock);
+	mutex_init(&edgetpu_pdev->platform_pwr.state_lock);
 	abrolhos_pwr_debugfs_dir =
 		debugfs_create_dir("power", edgetpu_fs_debugfs_dir());
 	if (!abrolhos_pwr_debugfs_dir) {
@@ -587,8 +648,10 @@ static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
 		/* don't fail the procedure on debug FS creation fails */
 		return 0;
 	}
-	debugfs_create_file("state", 0660, abrolhos_pwr_debugfs_dir, dev,
+	debugfs_create_file("state", 0660, abrolhos_pwr_debugfs_dir, etdev,
 			    &fops_tpu_pwr_state);
+	debugfs_create_file("min_state", 0660, abrolhos_pwr_debugfs_dir, etdev,
+			    &fops_tpu_min_pwr_state);
 	debugfs_create_file("vdd_tpu", 0660, abrolhos_pwr_debugfs_dir, dev,
 			    &fops_tpu_vdd_tpu);
 	debugfs_create_file("vdd_tpu_m", 0660, abrolhos_pwr_debugfs_dir, dev,
diff --git a/drivers/edgetpu/abrolhos-thermal.c b/drivers/edgetpu/abrolhos-thermal.c
index bfa28cf..cbb9abe 100644
--- a/drivers/edgetpu/abrolhos-thermal.c
+++ b/drivers/edgetpu/abrolhos-thermal.c
@@ -14,6 +14,7 @@
 #include <linux/slab.h>
 #include <linux/thermal.h>
 #include <linux/version.h>
+#include <linux/of.h>
 
 #include "abrolhos-firmware.h"
 #include "abrolhos-platform.h"
@@ -23,42 +24,19 @@
 #include "edgetpu-mmu.h"
 #include "edgetpu-thermal.h"
 
-static const unsigned long state_mapping[] = {
-	TPU_ACTIVE_OD,
-	TPU_ACTIVE_NOM,
-	TPU_ACTIVE_UD,
-	TPU_ACTIVE_SUD,
-	TPU_RETENTION_CLOCKS_SLOW,
-	TPU_SLEEP_CLOCKS_SLOW,
-	TPU_SLEEP_CLOCKS_OFF,
-	TPU_DEEP_SLEEP_CLOCKS_FAST,
-	TPU_DEEP_SLEEP_CLOCKS_SLOW,
-	TPU_DEEP_SLEEP_CLOCKS_OFF,
-	TPU_OFF,
-};
-
-/*
- * Sequence need to be kept to make power increasing
- * to make sure we always return the highest state.
- */
-static const struct edgetpu_state_pwr state_pwr_map[] = {
-	{ TPU_ACTIVE_OD, 198 },
-	{ TPU_ACTIVE_NOM, 165 },
-	{ TPU_ACTIVE_UD, 131 },
-	{ TPU_ACTIVE_SUD, 102 },
-	{ TPU_SLEEP_CLOCKS_SLOW, 66 },
-	{ TPU_SLEEP_CLOCKS_OFF, 66 },
-	{ TPU_RETENTION_CLOCKS_SLOW, 49 },
-	{ TPU_DEEP_SLEEP_CLOCKS_FAST, 43 },
-	{ TPU_DEEP_SLEEP_CLOCKS_SLOW, 6 },
-	{ TPU_DEEP_SLEEP_CLOCKS_OFF, 6 },
-	{ TPU_OFF, 0 },
-};
+#define MAX_NUM_TPU_STATES 10
+#define OF_DATA_NUM_MAX MAX_NUM_TPU_STATES * 2
+static struct edgetpu_state_pwr state_pwr_map[MAX_NUM_TPU_STATES] = {0};
 
 static int edgetpu_get_max_state(struct thermal_cooling_device *cdev,
 				 unsigned long *state)
 {
-	*state = ARRAY_SIZE(state_mapping) - 1;
+	struct edgetpu_thermal *thermal = cdev->devdata;
+
+	if (thermal->tpu_num_states <= 0)
+		return -ENOSYS;
+
+	*state = thermal->tpu_num_states - 1;
 	return 0;
 }
 
@@ -73,14 +51,14 @@ static int edgetpu_set_cur_state(struct thermal_cooling_device *cdev,
 	struct device *dev = cooling->dev;
 	unsigned long pwr_state;
 
-	if (state_original >= ARRAY_SIZE(state_mapping)) {
+	if (state_original >= cooling->tpu_num_states) {
 		dev_err(dev, "%s: invalid cooling state %lu\n", __func__,
 			state_original);
 		return -EINVAL;
 	}
 
 	mutex_lock(&cooling->lock);
-	pwr_state = state_mapping[state_original];
+	pwr_state = state_pwr_map[state_original].state;
 	if (state_original != cooling->cooling_state) {
 		/*
 		 * Cap the minimum state we request here.
@@ -116,7 +94,7 @@ static int edgetpu_get_cur_state(struct thermal_cooling_device *cdev,
 	struct edgetpu_thermal *cooling = cdev->devdata;
 
 	*state = cooling->cooling_state;
-	if (*state >= ARRAY_SIZE(state_mapping)) {
+	if (*state >= cooling->tpu_num_states) {
 		dev_warn(cooling->dev,
 			 "Unknown cooling state: %lu, resetting\n", *state);
 		mutex_lock(&cooling->lock);
@@ -138,17 +116,17 @@ static int edgetpu_get_cur_state(struct thermal_cooling_device *cdev,
 }
 
 static int edgetpu_state2power_internal(unsigned long state, u32 *power,
-					struct device *dev)
+					struct edgetpu_thermal *thermal)
 {
 	int i;
 
-	for (i = 0; i < ARRAY_SIZE(state_pwr_map); i++) {
+	for (i = 0; i < thermal->tpu_num_states; ++i) {
 		if (state == state_pwr_map[i].state) {
 			*power = state_pwr_map[i].power;
 			return 0;
 		}
 	}
-	dev_err(dev, "Unknown state req for: %lu\n", state);
+	dev_err(thermal->dev, "Unknown state req for: %lu\n", state);
 	*power = 0;
 	return -EINVAL;
 }
@@ -164,7 +142,7 @@ static int edgetpu_get_requested_power(struct thermal_cooling_device *cdev,
 
 	state_original = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 	return edgetpu_state2power_internal(state_original, power,
-					    cooling->dev);
+					    cooling);
 }
 
 static int edgetpu_state2power(struct thermal_cooling_device *cdev,
@@ -175,14 +153,14 @@ static int edgetpu_state2power(struct thermal_cooling_device *cdev,
 {
 	struct edgetpu_thermal *cooling = cdev->devdata;
 
-	if (state >= ARRAY_SIZE(state_mapping)) {
+	if (state >= cooling->tpu_num_states) {
 		dev_err(cooling->dev, "%s: invalid state: %lu\n", __func__,
 			state);
 		return -EINVAL;
 	}
 
-	return edgetpu_state2power_internal(state_mapping[state], power,
-					    cooling->dev);
+	return edgetpu_state2power_internal(state_pwr_map[state].state, power,
+					    cooling);
 }
 
 static int edgetpu_power2state(struct thermal_cooling_device *cdev,
@@ -191,18 +169,29 @@ static int edgetpu_power2state(struct thermal_cooling_device *cdev,
 #endif
 			       u32 power, unsigned long *state)
 {
-	int i;
-	struct edgetpu_thermal *cooling = cdev->devdata;
-
-	for (i = 0; i < ARRAY_SIZE(state_pwr_map); i++) {
-		if (power >= state_pwr_map[i].power) {
-			*state = i;
-			return 0;
+	int i, penultimate_throttle_state;
+	struct edgetpu_thermal *thermal = cdev->devdata;
+
+	*state = 0;
+	if (thermal->tpu_num_states < 2)
+		return thermal->tpu_num_states == 1 ? 0 : -ENOSYS;
+
+	penultimate_throttle_state = thermal->tpu_num_states - 2;
+	/*
+	 * argument "power" is the maximum allowed power consumption in mW as
+	 * defined by the PID control loop. Check for the first state that is
+	 * less than or equal to the current allowed power. state_pwr_map is
+	 * descending, so lowest power consumption is last value in the array
+	 * return lowest state even if it consumes more power than allowed as
+	 * not all platforms can handle throttling below an active state
+	 */
+	for (i = penultimate_throttle_state; i >= 0; --i) {
+		if (power < state_pwr_map[i].power) {
+			*state = i + 1;
+			break;
 		}
 	}
-
-	dev_err(cooling->dev, "No power2state mapping found: %d\n", power);
-	return -EINVAL;
+	return 0;
 }
 
 static struct thermal_cooling_device_ops edgetpu_cooling_ops = {
@@ -233,12 +222,56 @@ static void devm_tpu_thermal_release(struct device *dev, void *res)
 	tpu_thermal_exit(thermal);
 }
 
+static int tpu_thermal_parse_dvfs_table(struct edgetpu_thermal *thermal)
+{
+	int row_size, col_size, tbl_size, i;
+	int of_data_int_array[OF_DATA_NUM_MAX];
+
+	if (of_property_read_u32_array(thermal->dev->of_node,
+                                "tpu_dvfs_table_size", of_data_int_array, 2 ))
+		goto error;
+
+	row_size = of_data_int_array[0];
+	col_size = of_data_int_array[1];
+	tbl_size = row_size * col_size;
+	if (row_size > MAX_NUM_TPU_STATES) {
+		dev_err(thermal->dev, "too many TPU states\n");
+		goto error;
+	}
+
+	if (tbl_size > OF_DATA_NUM_MAX)
+		goto error;
+
+	if (of_property_read_u32_array(thermal->dev->of_node,
+                                "tpu_dvfs_table", of_data_int_array, tbl_size))
+		goto error;
+
+	thermal->tpu_num_states = row_size;
+	for (i = 0; i < row_size; ++i) {
+		int idx = col_size * i;
+		state_pwr_map[i].state = of_data_int_array[idx];
+		state_pwr_map[i].power = of_data_int_array[idx+1];
+	}
+
+	return 0;
+
+error:
+	dev_err(thermal->dev, "failed to parse DVFS table\n");
+	return -EINVAL;
+}
+
 static int
 tpu_thermal_cooling_register(struct edgetpu_thermal *thermal, char *type)
 {
 	struct device_node *cooling_node = NULL;
+	int err = 0;
 
 	thermal->op_data = NULL;
+	thermal->tpu_num_states = 0;
+
+	err = tpu_thermal_parse_dvfs_table(thermal);
+	if (err)
+		return err;
 
 	mutex_init(&thermal->lock);
 	cooling_node = of_find_node_by_name(NULL, "tpu-cooling");
diff --git a/drivers/edgetpu/abrolhos/config.h b/drivers/edgetpu/abrolhos/config.h
index fa52c01..f6397fd 100644
--- a/drivers/edgetpu/abrolhos/config.h
+++ b/drivers/edgetpu/abrolhos/config.h
@@ -58,7 +58,7 @@
 /*
  * Size of memory for FW accessible debug dump segments
  */
-#define EDGETPU_DEBUG_DUMP_MEM_SIZE 0x40000
+#define EDGETPU_DEBUG_DUMP_MEM_SIZE 0x4E0000
 
 #include "config-mailbox.h"
 #include "config-tpu-cpu.h"
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 9de2a75..891cde6 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -91,9 +91,19 @@ static void edgetpu_vma_close(struct vm_area_struct *vma)
 {
 	struct edgetpu_client *client = vma->vm_private_data;
 	enum edgetpu_wakelock_event evt = mmap_wakelock_event(vma->vm_pgoff);
+	struct edgetpu_dev *etdev = client->etdev;
 
 	if (evt != EDGETPU_WAKELOCK_EVENT_END)
 		edgetpu_wakelock_dec_event(client->wakelock, evt);
+
+	/* TODO(b/184613387): check whole VMA range instead of the start only */
+	if (vma->vm_pgoff == EDGETPU_MMAP_LOG_BUFFER_OFFSET >> PAGE_SHIFT)
+		edgetpu_munmap_telemetry_buffer(etdev, EDGETPU_TELEMETRY_LOG,
+						vma);
+	else if (vma->vm_pgoff ==
+		 EDGETPU_MMAP_TRACE_BUFFER_OFFSET >> PAGE_SHIFT)
+		edgetpu_munmap_telemetry_buffer(etdev, EDGETPU_TELEMETRY_TRACE,
+						vma);
 }
 
 static const struct vm_operations_struct edgetpu_vma_ops = {
@@ -329,6 +339,7 @@ struct edgetpu_client *edgetpu_client_add(struct edgetpu_dev *etdev)
 	mutex_init(&client->group_lock);
 	/* equivalent to edgetpu_client_get() */
 	refcount_set(&client->count, 1);
+	client->perdie_events = 0;
 	return client;
 }
 
@@ -348,8 +359,11 @@ void edgetpu_client_put(struct edgetpu_client *client)
 
 void edgetpu_client_remove(struct edgetpu_client *client)
 {
+	struct edgetpu_dev *etdev;
+
 	if (IS_ERR_OR_NULL(client))
 		return;
+	etdev = client->etdev;
 	/*
 	 * A quick check without holding client->group_lock.
 	 *
@@ -367,6 +381,15 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	 * happen.
 	 */
 	client->wakelock = NULL;
+
+	/* 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);
+	if (client->perdie_events &
+	    1 << perdie_event_id_to_num(EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE))
+		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_TRACE);
+
 	edgetpu_client_put(client);
 }
 
@@ -387,6 +410,47 @@ void edgetpu_unregister_irq(struct edgetpu_dev *etdev, int irq)
 	devm_free_irq(etdev->dev, irq, etdev);
 }
 
+int edgetpu_alloc_coherent(struct edgetpu_dev *etdev, size_t size,
+			   struct edgetpu_coherent_mem *mem,
+			   enum edgetpu_context_id context_id)
+{
+	const u32 flags = EDGETPU_MMU_DIE | EDGETPU_MMU_32 | EDGETPU_MMU_HOST;
+
+	mem->vaddr = dma_alloc_coherent(etdev->dev, size, &mem->dma_addr,
+					GFP_KERNEL);
+	if (!mem->vaddr)
+		return -ENOMEM;
+	edgetpu_x86_coherent_mem_init(mem);
+	mem->tpu_addr =
+		edgetpu_mmu_tpu_map(etdev, mem->dma_addr, size,
+				    DMA_BIDIRECTIONAL, context_id, flags);
+	if (!mem->tpu_addr) {
+		dma_free_coherent(etdev->dev, size, mem->vaddr, mem->dma_addr);
+		mem->vaddr = NULL;
+		return -EINVAL;
+	}
+	mem->size = size;
+	return 0;
+}
+
+void edgetpu_free_coherent(struct edgetpu_dev *etdev,
+			   struct edgetpu_coherent_mem *mem,
+			   enum edgetpu_context_id context_id)
+{
+	edgetpu_mmu_tpu_unmap(etdev, mem->tpu_addr, mem->size, context_id);
+	edgetpu_x86_coherent_mem_set_wb(mem);
+	dma_free_coherent(etdev->dev, mem->size, mem->vaddr, mem->dma_addr);
+	mem->vaddr = NULL;
+}
+
+void edgetpu_handle_firmware_crash(struct edgetpu_dev *etdev, u16 crash_type,
+				   u32 extra_info)
+{
+	etdev_err(etdev, "firmware crashed: %u 0x%x", crash_type, extra_info);
+	etdev->firmware_crash_count++;
+	edgetpu_fatal_error_notify(etdev);
+}
+
 int __init edgetpu_init(void)
 {
 	int ret;
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 2b2f2c8..94129e1 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -608,19 +608,22 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 			   const struct edgetpu_mailbox_attr *attr)
 {
 	static uint cur_workload_id;
-	int ret = -EINVAL;
+	int ret;
 	struct edgetpu_device_group *group;
 	struct edgetpu_iommu_domain *etdomain;
 
-	if (!edgetpu_mailbox_validate_attr(attr))
+	ret = edgetpu_mailbox_validate_attr(attr);
+	if (ret)
 		goto error;
 	/*
 	 * The client already belongs to a group.
 	 * It's safe not to take client->group_lock as
 	 * edgetpu_device_group_add() will fail if there is race.
 	 */
-	if (client->group)
+	if (client->group) {
+		ret = -EINVAL;
 		goto error;
+	}
 
 	group = kzalloc(sizeof(*group), GFP_KERNEL);
 	if (!group) {
@@ -1077,7 +1080,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 
 	etdev_dbg(etdev, "%s: hostaddr=0x%llx pages=%u dir=%x", __func__,
 		  host_addr, num_pages, dir);
-	pages = kcalloc(num_pages, sizeof(*pages), GFP_KERNEL);
+	/*
+	 * "num_pages" is decided from user-space arguments, don't show warnings
+	 * when facing malicious input.
+	 */
+	pages = kcalloc(num_pages, sizeof(*pages), GFP_KERNEL | __GFP_NOWARN);
 	if (!pages)
 		return ERR_PTR(-ENOMEM);
 
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 14cbbce..9612c97 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -157,14 +157,20 @@ static int edgetpu_ioctl_unset_eventfd(struct edgetpu_client *client,
 }
 
 static int
-edgetpu_ioctl_set_perdie_eventfd(struct edgetpu_dev *etdev,
+edgetpu_ioctl_set_perdie_eventfd(struct edgetpu_client *client,
 				 struct edgetpu_event_register __user *argp)
 {
+	struct edgetpu_dev *etdev = client->etdev;
 	struct edgetpu_event_register eventreg;
 
 	if (copy_from_user(&eventreg, argp, sizeof(eventreg)))
 		return -EFAULT;
 
+	if (perdie_event_id_to_num(eventreg.event_id) >=
+	    EDGETPU_NUM_PERDIE_EVENTS)
+		return -EINVAL;
+	client->perdie_events |= 1 << perdie_event_id_to_num(eventreg.event_id);
+
 	switch (eventreg.event_id) {
 	case EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE:
 		return edgetpu_telemetry_set_event(etdev, EDGETPU_TELEMETRY_LOG,
@@ -177,9 +183,15 @@ edgetpu_ioctl_set_perdie_eventfd(struct edgetpu_dev *etdev,
 	}
 }
 
-static int edgetpu_ioctl_unset_perdie_eventfd(struct edgetpu_dev *etdev,
+static int edgetpu_ioctl_unset_perdie_eventfd(struct edgetpu_client *client,
 					      uint event_id)
 {
+	struct edgetpu_dev *etdev = client->etdev;
+
+	if (perdie_event_id_to_num(event_id) >= EDGETPU_NUM_PERDIE_EVENTS)
+		return -EINVAL;
+	client->perdie_events &= ~(1 << perdie_event_id_to_num(event_id));
+
 	switch (event_id) {
 	case EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE:
 		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_LOG);
@@ -629,13 +641,13 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 		ret = edgetpu_ioctl_finalize_group(client);
 		break;
 	case EDGETPU_SET_PERDIE_EVENTFD:
-		ret = edgetpu_ioctl_set_perdie_eventfd(client->etdev, argp);
+		ret = edgetpu_ioctl_set_perdie_eventfd(client, argp);
 		break;
 	case EDGETPU_UNSET_EVENT:
 		ret = edgetpu_ioctl_unset_eventfd(client, arg);
 		break;
 	case EDGETPU_UNSET_PERDIE_EVENT:
-		ret = edgetpu_ioctl_unset_perdie_eventfd(client->etdev, arg);
+		ret = edgetpu_ioctl_unset_perdie_eventfd(client, arg);
 		break;
 	case EDGETPU_SYNC_BUFFER:
 		ret = edgetpu_ioctl_sync_buffer(client, argp);
@@ -883,6 +895,25 @@ static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 			    &statusregs_ops);
 }
 
+static ssize_t firmware_crash_count_show(
+		struct device *dev, struct device_attribute *attr,
+		char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	return scnprintf(buf, PAGE_SIZE, "%u\n", etdev->firmware_crash_count);
+}
+static DEVICE_ATTR_RO(firmware_crash_count);
+
+static struct attribute *edgetpu_dev_attrs[] = {
+	&dev_attr_firmware_crash_count.attr,
+	NULL,
+};
+
+static const struct attribute_group edgetpu_attr_group = {
+	.attrs = edgetpu_dev_attrs,
+};
+
 const struct file_operations edgetpu_fops = {
 	.owner = THIS_MODULE,
 	.llseek = no_llseek,
@@ -919,12 +950,16 @@ int edgetpu_fs_add(struct edgetpu_dev *etdev)
 		return ret;
 	}
 
+	ret = device_add_group(etdev->dev, &edgetpu_attr_group);
+	if (ret)
+		etdev_warn(etdev, "edgetpu attr group create failed: %d", ret);
 	edgetpu_fs_setup_debugfs(etdev);
 	return 0;
 }
 
 void edgetpu_fs_remove(struct edgetpu_dev *etdev)
 {
+	device_remove_group(etdev->dev, &edgetpu_attr_group);
 	device_destroy(edgetpu_class, etdev->devno);
 	cdev_del(&etdev->cdev);
 	debugfs_remove_recursive(etdev->d_entry);
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 511823c..9d4f10d 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -35,8 +35,6 @@
 #include "edgetpu-thermal.h"
 #include "edgetpu-usage-stats.h"
 
-#include "soc/google/bcl.h"
-
 #define etdev_err(etdev, fmt, ...) dev_err((etdev)->etcdev, fmt, ##__VA_ARGS__)
 #define etdev_warn(etdev, fmt, ...)                                            \
 	dev_warn((etdev)->etcdev, fmt, ##__VA_ARGS__)
@@ -103,6 +101,10 @@ struct edgetpu_p2p_csr_map;
 struct edgetpu_remote_dram_map;
 struct edgetpu_wakelock;
 
+#define EDGETPU_NUM_PERDIE_EVENTS	2
+#define perdie_event_id_to_num(event_id)				      \
+	(event_id - EDGETPU_PERDIE_EVENT_LOGS_AVAILABLE)
+
 struct edgetpu_client {
 	pid_t pid;
 	pid_t tgid;
@@ -130,6 +132,8 @@ struct edgetpu_client {
 	struct edgetpu_reg_window reg_window;
 	/* Per-client request to keep device active */
 	struct edgetpu_wakelock *wakelock;
+	/* Bit field of registered per die events */
+	u64 perdie_events;
 };
 
 struct edgetpu_mapping;
@@ -197,11 +201,14 @@ struct edgetpu_dev {
 	/* version read from the firmware binary file */
 	struct edgetpu_fw_version fw_version;
 	atomic_t job_count;	/* times joined to a device group */
+
+	/* counts of error events */
+	uint firmware_crash_count;
+
 	struct edgetpu_coherent_mem debug_dump_mem;	/* debug dump memory */
 	/* debug dump handlers */
 	edgetpu_debug_dump_handlers *debug_dump_handlers;
 	struct work_struct debug_dump_work;
-	struct gs101_bcl_dev *bcl_dev;
 };
 
 extern const struct file_operations edgetpu_fops;
@@ -295,10 +302,29 @@ edgetpu_x86_coherent_mem_set_wb(struct edgetpu_coherent_mem *mem)
 #endif
 }
 
+/*
+ * Attempt to allocate memory from the dma coherent memory using dma_alloc.
+ * Use this to allocate memory outside the instruction remap pool.
+ */
+int edgetpu_alloc_coherent(struct edgetpu_dev *etdev, size_t size,
+			   struct edgetpu_coherent_mem *mem,
+			   enum edgetpu_context_id context_id);
+/*
+ * Free memory allocated by the function above from the dma coherent memory.
+ */
+void edgetpu_free_coherent(struct edgetpu_dev *etdev,
+			   struct edgetpu_coherent_mem *mem,
+			   enum edgetpu_context_id context_id);
+
+
 /* External drivers can hook up to edgetpu driver using these calls. */
 int edgetpu_open(struct edgetpu_dev *etdev, struct file *file);
 long edgetpu_ioctl(struct file *file, uint cmd, ulong arg);
 
+/* Handle firmware crash event */
+void edgetpu_handle_firmware_crash(struct edgetpu_dev *etdev, u16 crash_type,
+				   u32 extra_info);
+
 /* Bus (Platform/PCI) <-> Core API */
 
 int __init edgetpu_init(void);
diff --git a/drivers/edgetpu/edgetpu-iremap-pool.c b/drivers/edgetpu/edgetpu-iremap-pool.c
index 13b9977..b23bcc6 100644
--- a/drivers/edgetpu/edgetpu-iremap-pool.c
+++ b/drivers/edgetpu/edgetpu-iremap-pool.c
@@ -78,29 +78,6 @@ void edgetpu_iremap_pool_destroy(struct edgetpu_dev *etdev)
 	etdev->iremap_pool = NULL;
 }
 
-static int edgetpu_alloc_coherent(struct edgetpu_dev *etdev, size_t size,
-				  struct edgetpu_coherent_mem *mem,
-				  enum edgetpu_context_id context_id)
-{
-	const u32 flags = EDGETPU_MMU_DIE | EDGETPU_MMU_32 | EDGETPU_MMU_HOST;
-
-	mem->vaddr = dma_alloc_coherent(etdev->dev, size, &mem->dma_addr,
-					GFP_KERNEL);
-	if (!mem->vaddr)
-		return -ENOMEM;
-	edgetpu_x86_coherent_mem_init(mem);
-	mem->tpu_addr =
-		edgetpu_mmu_tpu_map(etdev, mem->dma_addr, size,
-				    DMA_BIDIRECTIONAL, context_id, flags);
-	if (!mem->tpu_addr) {
-		dma_free_coherent(etdev->dev, size, mem->vaddr, mem->dma_addr);
-		mem->vaddr = NULL;
-		return -EINVAL;
-	}
-	mem->size = size;
-	return 0;
-}
-
 int edgetpu_iremap_alloc(struct edgetpu_dev *etdev, size_t size,
 			 struct edgetpu_coherent_mem *mem,
 			 enum edgetpu_context_id context_id)
@@ -123,22 +100,12 @@ int edgetpu_iremap_alloc(struct edgetpu_dev *etdev, size_t size,
 	mem->dma_addr = etmempool->base_dma_addr + offset;
 	mem->tpu_addr = etmempool->base_tpu_addr + offset;
 	mem->size = size;
-	etdev_dbg(etdev, "iremap_alloc @ %llx IOVA = %llx size = %zu",
-		   (u64)mem->vaddr, mem->dma_addr, size);
+	etdev_dbg(etdev, "%s @ %llx IOVA = %llx size = %zu",
+		   __func__, (u64)mem->vaddr, mem->dma_addr, size);
 	mutex_unlock(&etmempool->lock);
 	return 0;
 }
 
-static void edgetpu_free_coherent(struct edgetpu_dev *etdev,
-				  struct edgetpu_coherent_mem *mem,
-				  enum edgetpu_context_id context_id)
-{
-	edgetpu_mmu_tpu_unmap(etdev, mem->tpu_addr, mem->size, context_id);
-	edgetpu_x86_coherent_mem_set_wb(mem);
-	dma_free_coherent(etdev->dev, mem->size, mem->vaddr, mem->dma_addr);
-	mem->vaddr = NULL;
-}
-
 void edgetpu_iremap_free(struct edgetpu_dev *etdev,
 			 struct edgetpu_coherent_mem *mem,
 			 enum edgetpu_context_id context_id)
@@ -150,8 +117,8 @@ void edgetpu_iremap_free(struct edgetpu_dev *etdev,
 		return;
 	}
 	mutex_lock(&etmempool->lock);
-	etdev_dbg(etdev, "iremap_free @ %llx IOVA = %llx size = %zu",
-		  (u64)mem->vaddr, mem->dma_addr, mem->size);
+	etdev_dbg(etdev, "%s @ %llx IOVA = %llx size = %zu",
+		  __func__, (u64)mem->vaddr, mem->dma_addr, mem->size);
 	gen_pool_free(etmempool->gen_pool, (unsigned long)mem->vaddr,
 		      mem->size);
 	mem->vaddr = NULL;
@@ -188,8 +155,8 @@ int edgetpu_iremap_mmap(struct edgetpu_dev *etdev, struct vm_area_struct *vma,
 
 	offset = mem->vaddr - etmempool->base_vaddr;
 	phys = etmempool->base_phys_addr + offset;
-	etdev_dbg(etdev, "iremap_mmap: virt = %llx phys = %llx\n",
-		  (u64)mem->vaddr, phys);
+	etdev_dbg(etdev, "%s: virt = %llx phys = %llx\n",
+		  __func__, (u64)mem->vaddr, phys);
 	ret = remap_pfn_range(vma, vma->vm_start, phys >> PAGE_SHIFT,
 			      vma->vm_end - vma->vm_start, vma->vm_page_prot);
 	vma->vm_pgoff = orig_pgoff;
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index e56eb9a..d48d799 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -87,9 +87,16 @@ edgetpu_reverse_kci_consume_response(struct edgetpu_dev *etdev,
 		edgetpu_chip_handle_reverse_kci(etdev, resp);
 		return;
 	}
-	/* We don't have any generic reverse KCI codes yet */
-	etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n", __func__,
-		   resp->code);
+
+	switch (resp->code) {
+	case RKCI_FIRMWARE_CRASH:
+		edgetpu_handle_firmware_crash(etdev, resp->status,
+					      resp->retval);
+		break;
+	default:
+		etdev_warn(etdev, "%s: Unrecognized KCI request: 0x%x\n",
+			   __func__, resp->code);
+	}
 }
 
 /* Remove one element from the circular buffer */
@@ -143,10 +150,10 @@ static int edgetpu_reverse_kci_add_response(
 	const struct edgetpu_kci_response_element *resp)
 {
 	struct edgetpu_reverse_kci *rkci = &kci->rkci;
-	unsigned long head, tail;
+	unsigned long head, tail, flags;
 	int ret = 0;
 
-	spin_lock(&rkci->producer_lock);
+	spin_lock_irqsave(&rkci->producer_lock, flags);
 	head = rkci->head;
 	tail = READ_ONCE(rkci->tail);
 	if (CIRC_SPACE(head, tail, REVERSE_KCI_BUFFER_SIZE) >= 1) {
@@ -157,7 +164,7 @@ static int edgetpu_reverse_kci_add_response(
 	} else {
 		ret = -ENOSPC;
 	}
-	spin_unlock(&rkci->producer_lock);
+	spin_unlock_irqrestore(&rkci->producer_lock, flags);
 	return ret;
 }
 
@@ -192,8 +199,9 @@ static void edgetpu_kci_consume_wait_list(
 		const struct edgetpu_kci_response_element *resp)
 {
 	struct edgetpu_kci_wait_list *cur, *nxt;
+	unsigned long flags;
 
-	spin_lock(&kci->wait_list_lock);
+	spin_lock_irqsave(&kci->wait_list_lock, flags);
 
 	list_for_each_entry_safe(cur, nxt, &kci->wait_list, list) {
 		if (cur->resp->seq > resp->seq)
@@ -210,7 +218,7 @@ static void edgetpu_kci_consume_wait_list(
 		kfree(cur);
 	}
 
-	spin_unlock(&kci->wait_list_lock);
+	spin_unlock_irqrestore(&kci->wait_list_lock, flags);
 }
 
 /*
@@ -530,15 +538,16 @@ void edgetpu_kci_release(struct edgetpu_dev *etdev, struct edgetpu_kci *kci)
 static int edgetpu_kci_push_wait_resp(struct edgetpu_kci *kci,
 				      struct edgetpu_kci_response_element *resp)
 {
-	struct edgetpu_kci_wait_list *entry = kzalloc(sizeof(*entry),
-						      GFP_KERNEL);
+	struct edgetpu_kci_wait_list *entry =
+		kzalloc(sizeof(*entry), GFP_KERNEL);
+	unsigned long flags;
 
 	if (!entry)
 		return -ENOMEM;
 	entry->resp = resp;
-	spin_lock(&kci->wait_list_lock);
+	spin_lock_irqsave(&kci->wait_list_lock, flags);
 	list_add_tail(&entry->list, &kci->wait_list);
-	spin_unlock(&kci->wait_list_lock);
+	spin_unlock_irqrestore(&kci->wait_list_lock, flags);
 
 	return 0;
 }
@@ -552,8 +561,9 @@ static void edgetpu_kci_del_wait_resp(struct edgetpu_kci *kci,
 				      struct edgetpu_kci_response_element *resp)
 {
 	struct edgetpu_kci_wait_list *cur;
+	unsigned long flags;
 
-	spin_lock(&kci->wait_list_lock);
+	spin_lock_irqsave(&kci->wait_list_lock, flags);
 
 	list_for_each_entry(cur, &kci->wait_list, list) {
 		if (cur->resp->seq > resp->seq)
@@ -565,7 +575,7 @@ static void edgetpu_kci_del_wait_resp(struct edgetpu_kci *kci,
 		}
 	}
 
-	spin_unlock(&kci->wait_list_lock);
+	spin_unlock_irqrestore(&kci->wait_list_lock, flags);
 }
 
 int edgetpu_kci_push_cmd(struct edgetpu_kci *kci,
@@ -769,8 +779,8 @@ int edgetpu_kci_leave_group(struct edgetpu_kci *kci)
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
-enum edgetpu_fw_flavor edgetpu_kci_fw_info(
-	struct edgetpu_kci *kci, struct edgetpu_fw_info *fw_info)
+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 = {
@@ -781,10 +791,6 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 		},
 	};
 	struct edgetpu_coherent_mem mem;
-	/* TODO(b/136208139): remove when old fw no longer in use */
-	struct edgetpu_command_element cmd_compat = {
-		.code = KCI_CODE_FIRMWARE_FLAVOR_COMPAT,
-	};
 	struct edgetpu_kci_response_element resp;
 	enum edgetpu_fw_flavor flavor = FW_FLAVOR_UNKNOWN;
 	int ret;
@@ -804,9 +810,6 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 	}
 
 	ret = edgetpu_kci_send_cmd_return_resp(kci, &cmd, &resp);
-	if (ret == KCI_ERROR_UNIMPLEMENTED)
-		ret = edgetpu_kci_send_cmd_return_resp(kci, &cmd_compat, &resp);
-
 	if (cmd.dma.address) {
 		memcpy(fw_info, mem.vaddr, sizeof(*fw_info));
 		edgetpu_iremap_free(etdev, &mem, EDGETPU_CONTEXT_KCI);
@@ -815,16 +818,16 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 	if (ret == KCI_ERROR_UNIMPLEMENTED) {
 		etdev_dbg(etdev, "old firmware does not report flavor\n");
 	} else if (ret == KCI_ERROR_OK) {
-		switch (resp.retval) {
+		switch (fw_info->fw_flavor) {
 		case FW_FLAVOR_BL1:
 		case FW_FLAVOR_SYSTEST:
 		case FW_FLAVOR_PROD_DEFAULT:
 		case FW_FLAVOR_CUSTOM:
-			flavor = resp.retval;
+			flavor = fw_info->fw_flavor;
 			break;
 		default:
 			etdev_dbg(etdev, "unrecognized fw flavor 0x%x\n",
-				  resp.retval);
+				  fw_info->fw_flavor);
 		}
 	} else {
 		etdev_dbg(etdev, "firmware flavor query returns %d\n", ret);
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index d9e7aeb..98a61df 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -70,14 +70,15 @@ struct edgetpu_kci_response_element {
 	u64 seq;
 	u16 code;
 	/*
-	 * Reserved on firmware side - they can't touch this.
+	 * Reserved on firmware side for KCI response - they can't touch this.
 	 * If a value is written here it will be discarded and overwritten
-	 * during response processing
+	 * during response processing.
+	 * For a Reverse KCI command firmware does set this as value1.
 	 */
 	u16 status;
 	/*
-	 * Return value is not currently needed by KCI command responses, but
-	 * incoming requests from firmware may encode information here.
+	 * Return value is not currently needed by KCI command responses.
+	 * For reverse KCI commands this is set as value2.
 	 */
 	u32 retval;
 } __packed;
@@ -102,8 +103,6 @@ enum edgetpu_kci_code {
 	KCI_CODE_JOIN_GROUP = 3,
 	KCI_CODE_LEAVE_GROUP = 4,
 	KCI_CODE_MAP_TRACE_BUFFER = 5,
-	/* TODO(b/136208139): remove when old fw no longer in use */
-	KCI_CODE_FIRMWARE_FLAVOR_COMPAT = 6,
 	KCI_CODE_SHUTDOWN = 7,
 	KCI_CODE_GET_DEBUG_DUMP = 8,
 	KCI_CODE_OPEN_DEVICE = 9,
@@ -122,6 +121,7 @@ enum edgetpu_reverse_kci_code {
 	RKCI_CHIP_CODE_FIRST = 0,
 	RKCI_CHIP_CODE_LAST = 0x7FFF,
 	RKCI_GENERIC_CODE_FIRST = 0x8000,
+	RKCI_FIRMWARE_CRASH = RKCI_GENERIC_CODE_FIRST + 0,
 	RKCI_GENERIC_CODE_LAST = 0xFFFF,
 };
 
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 3bdc7c5..6ffc141 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -5,10 +5,12 @@
  * Copyright (C) 2019 Google, Inc.
  */
 
+#include <asm/page.h>
 #include <linux/bits.h>
 #include <linux/dma-mapping.h>
 #include <linux/err.h>
 #include <linux/kernel.h>
+#include <linux/mmzone.h> /* MAX_ORDER_NR_PAGES */
 #include <linux/slab.h>
 
 #include "edgetpu-device-group.h"
@@ -378,8 +380,8 @@ static int convert_runtime_queue_size_to_fw(u32 queue_size, u32 element_size)
 	/* zero size is not allowed */
 	if (queue_size == 0 || element_size == 0)
 		return -EINVAL;
-	/* prevent integer overflow */
-	if (queue_size > SIZE_MAX / runtime_unit)
+	/* A quick check to prevent the queue allocation failure. */
+	if (queue_size > (MAX_ORDER_NR_PAGES << PAGE_SHIFT) / runtime_unit)
 		return -ENOMEM;
 	/*
 	 * Kernel doesn't care whether queue_size * runtime_unit is a multiple
@@ -392,19 +394,19 @@ static int convert_runtime_queue_size_to_fw(u32 queue_size, u32 element_size)
 	return ret;
 }
 
-bool edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr)
+int edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr)
 {
 	int size;
 
 	size = convert_runtime_queue_size_to_fw(attr->cmd_queue_size,
 						attr->sizeof_cmd);
 	if (size < 0)
-		return false;
+		return size;
 	size = convert_runtime_queue_size_to_fw(attr->resp_queue_size,
 						attr->sizeof_resp);
 	if (size < 0)
-		return false;
-	return true;
+		return size;
+	return 0;
 }
 
 int edgetpu_mailbox_init_vii(struct edgetpu_vii *vii,
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index 026bdbf..ddae12a 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -216,12 +216,12 @@ struct edgetpu_mailbox *
 edgetpu_mailbox_vii_add(struct edgetpu_mailbox_manager *mgr, uint id);
 /*
  * Validates the mailbox attributes.
- * Returns true if valid, false otherwise.
+ * Returns 0 if valid, otherwise a negative errno.
  *
  * See the error cases of EDGETPU_CREATE_GROUP in edgetpu.h for when will @attr
  * be considered as invalid.
  */
-bool edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr);
+int edgetpu_mailbox_validate_attr(const struct edgetpu_mailbox_attr *attr);
 /*
  * Sets mailbox and allocates queues to @vii.
  *
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index 68351a8..0ade5b0 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -4,10 +4,12 @@
  *
  * Copyright (C) 2019-2020 Google, Inc.
  */
+
 #include <linux/delay.h>
 #include <linux/dma-mapping.h>
 #include <linux/errno.h>
 #include <linux/mm.h>
+#include <linux/mutex.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>
 #include <linux/workqueue.h>
@@ -74,11 +76,11 @@ static int telemetry_set_event(struct edgetpu_dev *etdev,
 	if (IS_ERR(ctx))
 		return PTR_ERR(ctx);
 
-	write_lock_irqsave(&tel->ctx_mem_lock, flags);
+	write_lock_irqsave(&tel->ctx_lock, flags);
 	if (tel->ctx)
 		eventfd_ctx_put(tel->ctx);
 	tel->ctx = ctx;
-	write_unlock_irqrestore(&tel->ctx_mem_lock, flags);
+	write_unlock_irqrestore(&tel->ctx_lock, flags);
 
 	return 0;
 }
@@ -90,11 +92,11 @@ static void telemetry_unset_event(struct edgetpu_dev *etdev,
 
 	if (!tel->inited)
 		return;
-	write_lock_irqsave(&tel->ctx_mem_lock, flags);
+	write_lock_irqsave(&tel->ctx_lock, flags);
 	if (tel->ctx)
 		eventfd_ctx_put(tel->ctx);
 	tel->ctx = NULL;
-	write_unlock_irqrestore(&tel->ctx_mem_lock, flags);
+	write_unlock_irqrestore(&tel->ctx_lock, flags);
 
 	return;
 }
@@ -205,12 +207,12 @@ static void telemetry_worker(struct work_struct *work)
 
 		prev_head = tel->header->head;
 		if (tel->header->head != tel->header->tail) {
-			read_lock(&tel->ctx_mem_lock);
+			read_lock(&tel->ctx_lock);
 			if (tel->ctx)
 				eventfd_signal(tel->ctx, 1);
 			else
 				tel->fallback_fn(tel);
-			read_unlock(&tel->ctx_mem_lock);
+			read_unlock(&tel->ctx_lock);
 		}
 
 		spin_unlock_irqrestore(&tel->state_lock, flags);
@@ -257,18 +259,37 @@ static int telemetry_mmap_buffer(struct edgetpu_dev *etdev,
 	if (!tel->inited)
 		return -ENODEV;
 
-	write_lock(&tel->ctx_mem_lock);
+	mutex_lock(&tel->mmap_lock);
 
-	ret = edgetpu_iremap_mmap(etdev, vma, &tel->coherent_mem);
+	if (!tel->is_mmapped) {
+		ret = edgetpu_iremap_mmap(etdev, vma, &tel->coherent_mem);
 
-	if (!ret)
-		tel->coherent_mem.host_addr = vma->vm_start;
+		if (!ret) {
+			tel->coherent_mem.host_addr = vma->vm_start;
+			tel->is_mmapped = true;
+		}
+	} else {
+		ret = -EBUSY;
+		etdev_warn(etdev, "Buffer is already mmapped");
+	}
 
-	write_unlock(&tel->ctx_mem_lock);
+	mutex_unlock(&tel->mmap_lock);
 
 	return ret;
 }
 
+static void telemetry_munmap_buffer(struct edgetpu_dev *etdev,
+				    struct edgetpu_telemetry *tel,
+				    struct vm_area_struct *vma)
+{
+	if (!tel->inited)
+		return;
+
+	mutex_lock(&tel->mmap_lock);
+	tel->is_mmapped = false;
+	mutex_unlock(&tel->mmap_lock);
+}
+
 static int telemetry_init(struct edgetpu_dev *etdev,
 			  struct edgetpu_telemetry *tel, const char *name,
 			  struct edgetpu_coherent_mem *mem,
@@ -306,7 +327,7 @@ static int telemetry_init(struct edgetpu_dev *etdev,
 		edgetpu_x86_coherent_mem_set_uc(&tel->coherent_mem);
 	}
 
-	rwlock_init(&tel->ctx_mem_lock);
+	rwlock_init(&tel->ctx_lock);
 	tel->name = name;
 	tel->etdev = etdev;
 
@@ -322,6 +343,8 @@ static int telemetry_init(struct edgetpu_dev *etdev,
 	tel->fallback_fn = fallback;
 	tel->state = EDGETPU_TELEMETRY_ENABLED;
 	tel->inited = true;
+	mutex_init(&tel->mmap_lock);
+	tel->is_mmapped = false;
 
 	return 0;
 }
@@ -443,3 +466,13 @@ int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev,
 	return telemetry_mmap_buffer(
 		etdev, select_telemetry(etdev->telemetry, type), vma);
 }
+
+void edgetpu_munmap_telemetry_buffer(struct edgetpu_dev *etdev,
+				     enum edgetpu_telemetry_type type,
+				     struct vm_area_struct *vma)
+{
+	if (!etdev->telemetry)
+		return;
+	telemetry_munmap_buffer(etdev, select_telemetry(etdev->telemetry, type),
+				vma);
+}
diff --git a/drivers/edgetpu/edgetpu-telemetry.h b/drivers/edgetpu/edgetpu-telemetry.h
index 308c8ea..0a57601 100644
--- a/drivers/edgetpu/edgetpu-telemetry.h
+++ b/drivers/edgetpu/edgetpu-telemetry.h
@@ -9,6 +9,7 @@
 
 #include <linux/eventfd.h>
 #include <linux/mm.h>
+#include <linux/mutex.h>
 #include <linux/seq_file.h>
 #include <linux/spinlock.h>
 #include <linux/types.h>
@@ -77,7 +78,7 @@ struct edgetpu_telemetry {
 	bool caller_mem;
 
 	struct eventfd_ctx *ctx; /* signal this to notify the runtime */
-	rwlock_t ctx_mem_lock; /* protects ctx and coherent_mem */
+	rwlock_t ctx_lock; /* protects ctx */
 	const char *name; /* for debugging */
 	bool inited; /* whether telemetry_init() succeeded */
 
@@ -85,6 +86,9 @@ struct edgetpu_telemetry {
 	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 is_mmapped */
+	/* Flag tracking when the telemetry buffer is mapped to user space. */
+	bool is_mmapped;
 };
 
 struct edgetpu_telemetry_ctx {
@@ -138,5 +142,8 @@ void edgetpu_telemetry_mappings_show(struct edgetpu_dev *etdev,
 int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev,
 				  enum edgetpu_telemetry_type type,
 				  struct vm_area_struct *vma);
+void edgetpu_munmap_telemetry_buffer(struct edgetpu_dev *etdev,
+				     enum edgetpu_telemetry_type type,
+				     struct vm_area_struct *vma);
 
 #endif /* __EDGETPU_TELEMETRY_H__ */
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index 4e97f07..7201597 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -21,6 +21,7 @@ struct edgetpu_thermal {
 	struct mutex lock;
 	void *op_data;
 	unsigned long cooling_state;
+	unsigned int tpu_num_states;
 };
 
 struct edgetpu_state_pwr {
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index 06328ae..421cd29 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -133,6 +133,24 @@ static void edgetpu_utilization_update(
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
+static void edgetpu_counter_update(
+	struct edgetpu_dev *etdev,
+	struct edgetpu_usage_counter *counter)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (!ustats)
+		return;
+
+	etdev_dbg(etdev, "%s: type=%d value=%llu\n", __func__,
+		  counter->type, counter->value);
+
+	mutex_lock(&ustats->usage_stats_lock);
+	if (counter->type >= 0 && counter->type < EDGETPU_COUNTER_COUNT)
+		ustats->counter[counter->type] = counter->value;
+	mutex_unlock(&ustats->usage_stats_lock);
+}
+
 void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 {
 	struct edgetpu_usage_header *header = buf;
@@ -157,6 +175,9 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 			edgetpu_utilization_update(
 				etdev, &metric->component_activity);
 			break;
+		case EDGETPU_METRIC_TYPE_COUNTER:
+			edgetpu_counter_update(etdev, &metric->counter);
+			break;
 		default:
 			etdev_dbg(etdev, "%s: %d: skip unknown type=%u",
 				  __func__, i, metric->type);
@@ -183,6 +204,22 @@ int edgetpu_usage_get_utilization(struct edgetpu_dev *etdev,
 	return val;
 }
 
+static int64_t edgetpu_usage_get_counter(
+	struct edgetpu_dev *etdev,
+	enum edgetpu_usage_counter_type counter_type)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int32_t val;
+
+	if (counter_type >= EDGETPU_COUNTER_COUNT)
+		return -1;
+	edgetpu_kci_update_usage(etdev);
+	mutex_lock(&ustats->usage_stats_lock);
+	val = ustats->counter[counter_type];
+	mutex_unlock(&ustats->usage_stats_lock);
+	return val;
+}
+
 static ssize_t tpu_usage_show(struct device *dev,
 			      struct device_attribute *attr,
 			      char *buf)
@@ -284,10 +321,38 @@ static ssize_t tpu_utilization_show(struct device *dev,
 }
 static DEVICE_ATTR_RO(tpu_utilization);
 
+static ssize_t tpu_active_cycle_count_show(struct device *dev,
+					   struct device_attribute *attr,
+					   char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	int64_t val;
+
+	val = edgetpu_usage_get_counter(etdev,
+					EDGETPU_COUNTER_TPU_ACTIVE_CYCLES);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(tpu_active_cycle_count);
+
+static ssize_t tpu_throttle_stall_count_show(struct device *dev,
+					     struct device_attribute *attr,
+					     char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	int64_t val;
+
+	val = edgetpu_usage_get_counter(etdev,
+					EDGETPU_COUNTER_TPU_THROTTLE_STALLS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(tpu_throttle_stall_count);
+
 static struct attribute *usage_stats_dev_attrs[] = {
 	&dev_attr_tpu_usage.attr,
 	&dev_attr_device_utilization.attr,
 	&dev_attr_tpu_utilization.attr,
+	&dev_attr_tpu_active_cycle_count.attr,
+	&dev_attr_tpu_throttle_stall_count.attr,
 	NULL,
 };
 
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 644d66d..a971ad7 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -53,11 +53,34 @@ struct edgetpu_component_activity {
 	int32_t utilization;
 };
 
+/*
+ * Defines different counter types we track.
+ * Must be kept in sync with firmware enum class CounterType.
+ */
+enum edgetpu_usage_counter_type {
+	/* TPU active cycles. */
+	EDGETPU_COUNTER_TPU_ACTIVE_CYCLES = 0,
+	/* Number of stalls caused by throttling. */
+	EDGETPU_COUNTER_TPU_THROTTLE_STALLS = 1,
+
+	EDGETPU_COUNTER_COUNT = 2, /* number of counters above */
+};
+
+/* Generic counter. Only reported if it has a value larger than 0. */
+struct __packed edgetpu_usage_counter {
+	/* What it counts. */
+	enum edgetpu_usage_counter_type type;
+
+	/* Accumulated value since last initialization. */
+	uint64_t value;
+};
+
 /* Must be kept in sync with firmware enum class UsageTrackerMetric::Type */
 enum edgetpu_usage_metric_type {
 	EDGETPU_METRIC_TYPE_RESERVED = 0,
 	EDGETPU_METRIC_TYPE_TPU_USAGE = 1,
 	EDGETPU_METRIC_TYPE_COMPONENT_ACTIVITY = 2,
+	EDGETPU_METRIC_TYPE_COUNTER = 3,
 };
 
 /*
@@ -70,6 +93,7 @@ struct edgetpu_usage_metric {
 	union {
 		struct tpu_usage tpu_usage;
 		struct edgetpu_component_activity component_activity;
+		struct edgetpu_usage_counter counter;
 	};
 };
 
@@ -79,6 +103,7 @@ struct edgetpu_usage_stats {
 	DECLARE_HASHTABLE(uid_hash_table, UID_HASH_BITS);
 	/* component utilization values reported by firmware */
 	int32_t component_utilization[EDGETPU_USAGE_COMPONENT_COUNT];
+	int64_t counter[EDGETPU_COUNTER_COUNT];
 	struct mutex usage_stats_lock;
 };