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

* whitechapel:
  edgetpu: Fix error handling for unpin_user_pages
  edgetpu: check DMA direction in map flags
  edgetpu: use entire page for sg table allocation
  edgetpu: disallow mmap buffer copies
  edgetpu: introduce edgetpu_vma_type to handle mmap
  edgetpu: check mailbox attached on mmap queues
  edgetpu: add size check on iremap mmap
  edgetpu: use num_pages for sg table allocation
  edgetpu: use max function for thread stats max stack values
  edgetpu: handle firmware thread stats
  edgetpu: make counters stats additive
  edgetpu: add firmware max watermark usage tracking
  edgetpu: abrolhos: implement firmware restart handler
  edgetpu: abrolhos: fix BTS scenario cleanup
  edgetpu: add additional counters to usage stats
  edgetpu: add firmware thread stats types to usage tracker

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

16 files changed, 757 insertions, 248 deletions

diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index 70520fa..8d08e33 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -85,18 +85,18 @@ u64 edgetpu_chip_tpu_timestamp(struct edgetpu_dev *etdev)
 void edgetpu_chip_init(struct edgetpu_dev *etdev)
 {
 	int i;
-	struct abrolhos_platform_dev *etpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 
 	/* Disable the CustomBlock Interrupt. */
 	edgetpu_dev_write_32(etdev, HOST_NONSECURE_INTRSRCMASKREG, 0x1);
 
-	if (!etpdev->ssmt_base)
+	if (!abpdev->ssmt_base)
 		return;
 
 	/* Setup non-secure SCIDs, assume VID = SCID */
 	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
-		writel(i, SSMT_NS_READ_STREAM_VID_REG(etpdev->ssmt_base, i));
-		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(etpdev->ssmt_base, i));
+		writel(i, SSMT_NS_READ_STREAM_VID_REG(abpdev->ssmt_base, i));
+		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(abpdev->ssmt_base, i));
 	}
 }
 
diff --git a/drivers/edgetpu/abrolhos-firmware.c b/drivers/edgetpu/abrolhos-firmware.c
index 4170c49..5fbec26 100644
--- a/drivers/edgetpu/abrolhos-firmware.c
+++ b/drivers/edgetpu/abrolhos-firmware.c
@@ -22,11 +22,11 @@ static int abrolhos_firmware_alloc_buffer(
 		struct edgetpu_firmware_buffer *fw_buf)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 
 	/* Allocate extra space the image header */
 	size_t buffer_size =
-		edgetpu_pdev->fw_region_size + MOBILE_FW_HEADER_SIZE;
+		abpdev->fw_region_size + MOBILE_FW_HEADER_SIZE;
 
 	fw_buf->vaddr = kzalloc(buffer_size, GFP_KERNEL);
 	if (!fw_buf->vaddr) {
@@ -64,11 +64,30 @@ static void abrolhos_firmware_teardown_buffer(
 {
 }
 
+static int abrolhos_firmware_restart(struct edgetpu_firmware *et_fw)
+{
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	int tpu_state;
+
+	tpu_state = gsa_send_tpu_cmd(edgetpu_pdev->gsa_dev, GSA_TPU_START);
+
+	if (tpu_state < 0) {
+		etdev_err(etdev, "GSA restart firmware failed: %d\n",
+			  tpu_state);
+		return -EIO;
+	}
+
+	etdev_dbg(etdev, "Firmware restart successful\n");
+
+	return 0;
+}
+
 static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 					 struct edgetpu_firmware_buffer *fw_buf)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 	void *image_vaddr, *header_vaddr;
 	struct mobile_image_config *image_config;
 	phys_addr_t image_start, image_end, carveout_start, carveout_end;
@@ -81,7 +100,7 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 		return -EINVAL;
 	}
 
-	tpu_state = gsa_send_tpu_cmd(edgetpu_pdev->gsa_dev, GSA_TPU_GET_STATE);
+	tpu_state = gsa_send_tpu_cmd(abpdev->gsa_dev, GSA_TPU_GET_STATE);
 
 	if (tpu_state < GSA_TPU_STATE_INACTIVE) {
 		etdev_warn(etdev, "GSA failed to retrieve current status: %d\n",
@@ -93,15 +112,15 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	etdev_dbg(etdev, "GSA Reports TPU state: %d\n", tpu_state);
 
 	if (tpu_state > GSA_TPU_STATE_INACTIVE) {
-		ret = gsa_unload_tpu_fw_image(edgetpu_pdev->gsa_dev);
+		ret = gsa_unload_tpu_fw_image(abpdev->gsa_dev);
 		if (ret) {
 			etdev_warn(etdev, "GSA release failed: %d\n", ret);
 			return -EIO;
 		}
 	}
 
-	image_vaddr = memremap(edgetpu_pdev->fw_region_paddr,
-			       edgetpu_pdev->fw_region_size, MEMREMAP_WC);
+	image_vaddr = memremap(abpdev->fw_region_paddr, abpdev->fw_region_size,
+			       MEMREMAP_WC);
 
 	if (!image_vaddr) {
 		etdev_err(etdev, "memremap failed\n");
@@ -113,7 +132,7 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
 
 	/* Allocate coherent memory for the image header */
-	header_vaddr = dma_alloc_coherent(edgetpu_pdev->gsa_dev,
+	header_vaddr = dma_alloc_coherent(abpdev->gsa_dev,
 					  MOBILE_FW_HEADER_SIZE,
 					  &header_dma_addr, GFP_KERNEL);
 	if (!header_vaddr) {
@@ -126,10 +145,10 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	memcpy(header_vaddr, fw_buf->vaddr, MOBILE_FW_HEADER_SIZE);
 	etdev_dbg(etdev,
 		  "Requesting GSA image load. meta = %llX payload = %llX",
-		  header_dma_addr, (u64)edgetpu_pdev->fw_region_paddr);
+		  header_dma_addr, (u64)abpdev->fw_region_paddr);
 
-	ret = gsa_load_tpu_fw_image(edgetpu_pdev->gsa_dev, header_dma_addr,
-				    edgetpu_pdev->fw_region_paddr);
+	ret = gsa_load_tpu_fw_image(abpdev->gsa_dev, header_dma_addr,
+				    abpdev->fw_region_paddr);
 	if (ret) {
 		etdev_err(etdev, "GSA authentication failed: %d\n", ret);
 		ret = -EIO;
@@ -148,8 +167,8 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	image_start = (phys_addr_t)image_config->carveout_base;
 	image_end = (phys_addr_t)(image_config->firmware_base +
 				  image_config->firmware_size - 1);
-	carveout_start = edgetpu_pdev->fw_region_paddr;
-	carveout_end = carveout_start + edgetpu_pdev->fw_region_size - 1;
+	carveout_start = abpdev->fw_region_paddr;
+	carveout_end = carveout_start + abpdev->fw_region_size - 1;
 
 	/* Image must fit within the carveout */
 	if (image_start < carveout_start || image_end > carveout_end) {
@@ -165,7 +184,7 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
 	edgetpu_mailbox_reset(etdev->kci->mailbox);
 
-	tpu_state = gsa_send_tpu_cmd(edgetpu_pdev->gsa_dev, GSA_TPU_START);
+	tpu_state = gsa_send_tpu_cmd(abpdev->gsa_dev, GSA_TPU_START);
 
 	if (tpu_state < 0) {
 		etdev_err(etdev, "GSA start firmware failed: %d\n", tpu_state);
@@ -173,7 +192,7 @@ static int abrolhos_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	}
 
 out_free_gsa:
-	dma_free_coherent(edgetpu_pdev->gsa_dev, MOBILE_FW_HEADER_SIZE,
+	dma_free_coherent(abpdev->gsa_dev, MOBILE_FW_HEADER_SIZE,
 			  header_vaddr, header_dma_addr);
 out_unmap:
 	memunmap(image_vaddr);
@@ -186,6 +205,7 @@ static const struct edgetpu_firmware_handlers abrolhos_firmware_handlers = {
 	.setup_buffer = abrolhos_firmware_setup_buffer,
 	.teardown_buffer = abrolhos_firmware_teardown_buffer,
 	.prepare_run = abrolhos_firmware_prepare_run,
+	.restart = abrolhos_firmware_restart,
 };
 
 int mobile_edgetpu_firmware_create(struct edgetpu_dev *etdev)
diff --git a/drivers/edgetpu/abrolhos-platform.c b/drivers/edgetpu/abrolhos-platform.c
index cc7ea13..2661853 100644
--- a/drivers/edgetpu/abrolhos-platform.c
+++ b/drivers/edgetpu/abrolhos-platform.c
@@ -209,17 +209,17 @@ static int abrolhos_parse_ssmt(struct abrolhos_platform_dev *etpdev)
 static int edgetpu_platform_probe(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
-	struct abrolhos_platform_dev *edgetpu_pdev;
+	struct abrolhos_platform_dev *abpdev;
 	struct resource *r;
 	struct edgetpu_mapped_resource regs;
 	int ret;
 
-	edgetpu_pdev = devm_kzalloc(dev, sizeof(*edgetpu_pdev), GFP_KERNEL);
-	if (!edgetpu_pdev)
+	abpdev = devm_kzalloc(dev, sizeof(*abpdev), GFP_KERNEL);
+	if (!abpdev)
 		return -ENOMEM;
 
-	platform_set_drvdata(pdev, &edgetpu_pdev->edgetpu_dev);
-	edgetpu_pdev->edgetpu_dev.dev = dev;
+	platform_set_drvdata(pdev, &abpdev->edgetpu_dev);
+	abpdev->edgetpu_dev.dev = dev;
 
 	r = platform_get_resource(pdev, IORESOURCE_MEM, 0);
 	if (IS_ERR_OR_NULL(r)) {
@@ -235,17 +235,17 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 		return -ENODEV;
 	}
 
-	mutex_init(&edgetpu_pdev->platform_pwr.policy_lock);
-	edgetpu_pdev->platform_pwr.curr_policy = TPU_POLICY_MAX;
+	mutex_init(&abpdev->platform_pwr.policy_lock);
+	abpdev->platform_pwr.curr_policy = TPU_POLICY_MAX;
 
-	ret = abrolhos_pm_create(&edgetpu_pdev->edgetpu_dev);
+	ret = abrolhos_pm_create(&abpdev->edgetpu_dev);
 
 	if (ret) {
 		dev_err(dev, "Failed to initialize PM interface (%d)\n", ret);
 		return ret;
 	}
 
-	ret = edgetpu_platform_setup_fw_region(edgetpu_pdev);
+	ret = edgetpu_platform_setup_fw_region(abpdev);
 	if (ret) {
 		dev_err(dev, "%s setup fw regions failed: %d\n", DRIVER_NAME,
 			ret);
@@ -253,15 +253,15 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 	}
 
 	ret = edgetpu_iremap_pool_create(
-		&edgetpu_pdev->edgetpu_dev,
+		&abpdev->edgetpu_dev,
 		/* Base virtual address (kernel address space) */
-		edgetpu_pdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
+		abpdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
 		/* Base DMA address */
 		EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
 		/* Base TPU address */
 		EDGETPU_REMAPPED_DATA_ADDR + EDGETPU_POOL_MEM_OFFSET,
 		/* Base physical address */
-		edgetpu_pdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
+		abpdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
 		/* Size */
 		EDGETPU_REMAPPED_DATA_SIZE - EDGETPU_POOL_MEM_OFFSET,
 		/* Granularity */
@@ -273,14 +273,13 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 		goto out_cleanup_fw;
 	}
 
-	edgetpu_pdev->edgetpu_dev.mcp_id = -1;
-	edgetpu_pdev->edgetpu_dev.mcp_die_index = 0;
-	edgetpu_pdev->irq = platform_get_irq(pdev, 0);
-	ret = edgetpu_device_add(&edgetpu_pdev->edgetpu_dev, &regs);
+	abpdev->edgetpu_dev.mcp_id = -1;
+	abpdev->edgetpu_dev.mcp_die_index = 0;
+	abpdev->irq = platform_get_irq(pdev, 0);
+	ret = edgetpu_device_add(&abpdev->edgetpu_dev, &regs);
 
-	if (!ret && edgetpu_pdev->irq >= 0)
-		ret = edgetpu_register_irq(&edgetpu_pdev->edgetpu_dev,
-					   edgetpu_pdev->irq);
+	if (!ret && abpdev->irq >= 0)
+		ret = edgetpu_register_irq(&abpdev->edgetpu_dev, abpdev->irq);
 
 	if (ret) {
 		dev_err(dev, "%s edgetpu setup failed: %d\n", DRIVER_NAME,
@@ -288,25 +287,24 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 		goto out_destroy_iremap;
 	}
 
-	ret = abrolhos_parse_ssmt(edgetpu_pdev);
+	ret = abrolhos_parse_ssmt(abpdev);
 	if (ret)
 		dev_warn(
 			dev,
 			"SSMT setup failed (%d). Context isolation not enforced\n",
 			ret);
 
-	abrolhos_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_LOG,
-				   &edgetpu_pdev->log_mem);
-	abrolhos_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_TRACE,
-				   &edgetpu_pdev->trace_mem);
+	abrolhos_get_telemetry_mem(abpdev, EDGETPU_TELEMETRY_LOG,
+				   &abpdev->log_mem);
+	abrolhos_get_telemetry_mem(abpdev, EDGETPU_TELEMETRY_TRACE,
+				   &abpdev->trace_mem);
 
-	ret = edgetpu_telemetry_init(&edgetpu_pdev->edgetpu_dev,
-				     &edgetpu_pdev->log_mem,
-				     &edgetpu_pdev->trace_mem);
+	ret = edgetpu_telemetry_init(&abpdev->edgetpu_dev, &abpdev->log_mem,
+				     &abpdev->trace_mem);
 	if (ret)
 		goto out_remove_device;
 
-	ret = mobile_edgetpu_firmware_create(&edgetpu_pdev->edgetpu_dev);
+	ret = mobile_edgetpu_firmware_create(&abpdev->edgetpu_dev);
 	if (ret) {
 		dev_err(dev,
 			"%s initialize firmware downloader failed: %d\n",
@@ -315,47 +313,47 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 	}
 
 	dev_dbg(dev, "Creating thermal device\n");
-	edgetpu_pdev->edgetpu_dev.thermal = devm_tpu_thermal_create(dev);
+	abpdev->edgetpu_dev.thermal = devm_tpu_thermal_create(dev);
 
 	dev_info(dev, "%s edgetpu initialized. Build: %s\n",
-		 edgetpu_pdev->edgetpu_dev.dev_name, GIT_REPO_TAG);
+		 abpdev->edgetpu_dev.dev_name, GIT_REPO_TAG);
 
 	dev_dbg(dev, "Probe finished, powering down\n");
 	/* Turn the device off unless a client request is already received. */
-	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, false);
+	edgetpu_pm_shutdown(&abpdev->edgetpu_dev, false);
 
-	edgetpu_pdev->sscd_info.pdata = &sscd_pdata;
-	edgetpu_pdev->sscd_info.dev = &sscd_dev;
+	abpdev->sscd_info.pdata = &sscd_pdata;
+	abpdev->sscd_info.dev = &sscd_dev;
 
 	return ret;
 out_tel_exit:
-	edgetpu_telemetry_exit(&edgetpu_pdev->edgetpu_dev);
+	edgetpu_telemetry_exit(&abpdev->edgetpu_dev);
 out_remove_device:
-	edgetpu_device_remove(&edgetpu_pdev->edgetpu_dev);
+	edgetpu_device_remove(&abpdev->edgetpu_dev);
 out_destroy_iremap:
-	edgetpu_iremap_pool_destroy(&edgetpu_pdev->edgetpu_dev);
+	edgetpu_iremap_pool_destroy(&abpdev->edgetpu_dev);
 out_cleanup_fw:
-	edgetpu_platform_cleanup_fw_region(edgetpu_pdev);
+	edgetpu_platform_cleanup_fw_region(abpdev);
 out_shutdown:
 	dev_dbg(dev, "Probe finished with error %d, powering down\n", ret);
-	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, true);
+	edgetpu_pm_shutdown(&abpdev->edgetpu_dev, true);
 	return ret;
 }
 
 static int edgetpu_platform_remove(struct platform_device *pdev)
 {
 	struct edgetpu_dev *etdev = platform_get_drvdata(pdev);
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 
 	mobile_edgetpu_firmware_destroy(etdev);
-	if (edgetpu_pdev->irq >= 0)
-		edgetpu_unregister_irq(etdev, edgetpu_pdev->irq);
+	if (abpdev->irq >= 0)
+		edgetpu_unregister_irq(etdev, abpdev->irq);
 
 	edgetpu_pm_get(etdev->pm);
 	edgetpu_telemetry_exit(etdev);
 	edgetpu_device_remove(etdev);
 	edgetpu_iremap_pool_destroy(etdev);
-	edgetpu_platform_cleanup_fw_region(edgetpu_pdev);
+	edgetpu_platform_cleanup_fw_region(abpdev);
 	edgetpu_pm_put(etdev->pm);
 	edgetpu_pm_shutdown(etdev, true);
 	abrolhos_pm_destroy(etdev);
diff --git a/drivers/edgetpu/abrolhos-platform.h b/drivers/edgetpu/abrolhos-platform.h
index 980dc8e..84e155e 100644
--- a/drivers/edgetpu/abrolhos-platform.h
+++ b/drivers/edgetpu/abrolhos-platform.h
@@ -10,8 +10,11 @@
 #include <linux/device.h>
 #include <linux/io.h>
 #include <linux/kernel.h>
+#include <linux/mutex.h>
 #include <linux/types.h>
 #include <soc/google/bcl.h>
+#include <soc/google/bts.h>
+#include <soc/google/exynos_pm_qos.h>
 
 #include "abrolhos-debug-dump.h"
 #include "abrolhos-pm.h"
@@ -20,17 +23,24 @@
 #define to_abrolhos_dev(etdev)                                                 \
 	container_of(etdev, struct abrolhos_platform_dev, edgetpu_dev)
 
-struct edgetpu_platform_pwr {
+struct abrolhos_platform_pwr {
 	struct mutex policy_lock;
 	enum tpu_pwr_state curr_policy;
 	struct mutex state_lock;
 	u64 min_state;
 	u64 requested_state;
+	/* INT/MIF requests for memory bandwidth */
+	struct exynos_pm_qos_request int_min;
+	struct exynos_pm_qos_request mif_min;
+	/* BTS */
+	unsigned int performance_scenario;
+	int scenario_count;
+	struct mutex scenario_lock;
 };
 
 struct abrolhos_platform_dev {
 	struct edgetpu_dev edgetpu_dev;
-	struct edgetpu_platform_pwr platform_pwr;
+	struct abrolhos_platform_pwr platform_pwr;
 	int irq;
 	phys_addr_t fw_region_paddr;
 	size_t fw_region_size;
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 47fa7dc..0562f30 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -34,14 +34,6 @@
 #define PM_QOS_MIF_MASK			(0xFFFF)
 #define PM_QOS_FACTOR			(1000)
 
-/* INT/MIF requests for memory bandwidth */
-static struct exynos_pm_qos_request int_min;
-static struct exynos_pm_qos_request mif_min;
-
-/* BTS */
-static unsigned int performance_scenario;
-static atomic64_t scenario_count = ATOMIC_INIT(0);
-
 /* Default power state: the lowest power state that keeps firmware running */
 static int power_state = TPU_DEEP_SLEEP_CLOCKS_SLOW;
 
@@ -130,8 +122,8 @@ static int abrolhos_pwr_state_get_locked(void *data, u64 *val)
 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;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	int ret = 0;
 
 	mutex_lock(&platform_pwr->state_lock);
@@ -145,8 +137,8 @@ static int abrolhos_pwr_state_set(void *data, u64 val)
 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;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	int ret;
 
 	mutex_lock(&platform_pwr->state_lock);
@@ -158,8 +150,8 @@ static int abrolhos_pwr_state_get(void *data, u64 *val)
 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;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	int ret = 0;
 
 	mutex_lock(&platform_pwr->state_lock);
@@ -173,8 +165,8 @@ static int abrolhos_min_pwr_state_set(void *data, u64 val)
 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;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 
 	mutex_lock(&platform_pwr->state_lock);
 	*val = platform_pwr->min_state;
@@ -184,15 +176,15 @@ static int abrolhos_min_pwr_state_get(void *data, u64 *val)
 
 static int abrolhos_pwr_policy_set(void *data, u64 val)
 {
-	struct abrolhos_platform_dev *edgetpu_pdev = (typeof(edgetpu_pdev))data;
-	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	struct abrolhos_platform_dev *abpdev = (typeof(abpdev))data;
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	int ret;
 
 	mutex_lock(&platform_pwr->policy_lock);
 	ret = exynos_acpm_set_policy(TPU_ACPM_DOMAIN, val);
 
 	if (ret) {
-		dev_err(edgetpu_pdev->edgetpu_dev.dev,
+		dev_err(abpdev->edgetpu_dev.dev,
 			"unable to set policy %lld (ret %d)\n", val, ret);
 		mutex_unlock(&platform_pwr->policy_lock);
 		return ret;
@@ -205,8 +197,8 @@ static int abrolhos_pwr_policy_set(void *data, u64 val)
 
 static int abrolhos_pwr_policy_get(void *data, u64 *val)
 {
-	struct abrolhos_platform_dev *edgetpu_pdev = (typeof(edgetpu_pdev))data;
-	struct edgetpu_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	struct abrolhos_platform_dev *abpdev = (typeof(abpdev))data;
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 
 	mutex_lock(&platform_pwr->policy_lock);
 	*val = platform_pwr->curr_policy;
@@ -443,7 +435,7 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm);
 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 abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 	int ret = abrolhos_pwr_state_set(
 		etpm->etdev, abrolhos_get_initial_pwr_state(etdev->dev));
 	enum edgetpu_firmware_status firmware_status;
@@ -454,9 +446,9 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 		return ret;
 
 	/* Clear out log / trace buffers */
-	memset(edgetpu_pdev->log_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
+	memset(abpdev->log_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
-	memset(edgetpu_pdev->trace_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
+	memset(abpdev->trace_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
 #endif
 
 	edgetpu_chip_init(etdev);
@@ -509,10 +501,10 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 	if (ret) {
 		abrolhos_power_down(etpm);
 	} else {
-		if (!edgetpu_pdev->bcl_dev)
-			edgetpu_pdev->bcl_dev = gs101_retrieve_bcl_handle();
-		if (edgetpu_pdev->bcl_dev)
-			gs101_init_tpu_ratio(edgetpu_pdev->bcl_dev);
+		if (!abpdev->bcl_dev)
+			abpdev->bcl_dev = gs101_retrieve_bcl_handle();
+		if (abpdev->bcl_dev)
+			gs101_init_tpu_ratio(abpdev->bcl_dev);
 	}
 
 	return ret;
@@ -521,7 +513,7 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 static void
 abrolhos_pm_shutdown_firmware(struct abrolhos_platform_dev *etpdev,
 			      struct edgetpu_dev *etdev,
-			      struct abrolhos_platform_dev *edgetpu_pdev)
+			      struct abrolhos_platform_dev *abpdev)
 {
 	if (!edgetpu_pchannel_power_down(etdev, false))
 		return;
@@ -537,7 +529,7 @@ abrolhos_pm_shutdown_firmware(struct abrolhos_platform_dev *etpdev,
 	etdev_warn(etdev, "Forcing shutdown through power policy\n");
 	/* Request GSA shutdown to make sure the R52 core is reset */
 	gsa_send_tpu_cmd(etpdev->gsa_dev, GSA_TPU_SHUTDOWN);
-	abrolhos_pwr_policy_set(edgetpu_pdev, TPU_OFF);
+	abrolhos_pwr_policy_set(abpdev, TPU_OFF);
 	pm_runtime_put_sync(etdev->dev);
 	/*
 	 * TODO: experiment on hardware to verify if this delay
@@ -547,34 +539,43 @@ abrolhos_pm_shutdown_firmware(struct abrolhos_platform_dev *etpdev,
 	 */
 	msleep(100);
 	pm_runtime_get_sync(etdev->dev);
-	abrolhos_pwr_policy_set(edgetpu_pdev, TPU_ACTIVE_OD);
+	abrolhos_pwr_policy_set(abpdev, TPU_ACTIVE_OD);
 }
 
 static void abrolhos_pm_cleanup_bts_scenario(struct edgetpu_dev *etdev)
 {
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
 	if (!performance_scenario)
 		return;
-	while (atomic64_fetch_dec(&scenario_count) > 0) {
+
+	mutex_lock(&platform_pwr->scenario_lock);
+	while (platform_pwr->scenario_count) {
 		int ret = bts_del_scenario(performance_scenario);
 
 		if (ret) {
-			atomic64_set(&scenario_count, 0);
+			platform_pwr->scenario_count = 0;
 			etdev_warn_once(
 				etdev,
 				"error %d in cleaning up BTS scenario %u\n",
 				ret, performance_scenario);
-			return;
+			break;
 		}
+		platform_pwr->scenario_count--;
 	}
+	mutex_unlock(&platform_pwr->scenario_lock);
 }
 
 static void abrolhos_power_down(struct edgetpu_pm *etpm)
 {
 	struct edgetpu_dev *etdev = etpm->etdev;
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	u64 val;
 	int res;
-	int min_state = edgetpu_pdev->platform_pwr.min_state;
+	int min_state = platform_pwr->min_state;
 
 	etdev_info(etdev, "Powering down\n");
 
@@ -584,12 +585,6 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm)
 		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, &val)) {
 		etdev_warn(etdev, "Failed to read current power state\n");
 		val = TPU_ACTIVE_NOM;
@@ -602,30 +597,49 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm)
 	if (etdev->kci && edgetpu_firmware_status_locked(etdev) == FW_VALID) {
 		/* Update usage stats before we power off fw. */
 		edgetpu_kci_update_usage_locked(etdev);
-		abrolhos_pm_shutdown_firmware(edgetpu_pdev, etdev,
-					      edgetpu_pdev);
+		abrolhos_pm_shutdown_firmware(abpdev, etdev, abpdev);
 		edgetpu_kci_cancel_work_queues(etdev->kci);
 	}
 
-	res = gsa_send_tpu_cmd(edgetpu_pdev->gsa_dev, GSA_TPU_SHUTDOWN);
+	res = gsa_send_tpu_cmd(abpdev->gsa_dev, GSA_TPU_SHUTDOWN);
 	if (res < 0)
 		etdev_warn(etdev, "GSA shutdown request failed (%d)\n", res);
 	abrolhos_pwr_state_set(etdev, TPU_OFF);
+
+	/* Remove our vote for INT/MIF state (if any) */
+	exynos_pm_qos_update_request(&platform_pwr->int_min, 0);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, 0);
+
+	abrolhos_pm_cleanup_bts_scenario(etdev);
 }
 
 static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
 {
 	int ret;
 	struct edgetpu_dev *etdev = etpm->etdev;
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 	struct device *dev = etdev->dev;
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 
 	ret = abrolhos_pwr_state_init(dev);
 	if (ret)
 		return ret;
 
-	mutex_init(&edgetpu_pdev->platform_pwr.policy_lock);
-	mutex_init(&edgetpu_pdev->platform_pwr.state_lock);
+	mutex_init(&platform_pwr->policy_lock);
+	mutex_init(&platform_pwr->state_lock);
+	mutex_init(&platform_pwr->scenario_lock);
+
+	exynos_pm_qos_add_request(&platform_pwr->int_min,
+				  PM_QOS_DEVICE_THROUGHPUT, 0);
+	exynos_pm_qos_add_request(&platform_pwr->mif_min, PM_QOS_BUS_THROUGHPUT,
+				  0);
+
+	platform_pwr->performance_scenario =
+		bts_get_scenindex("tpu_performance");
+	if (!platform_pwr->performance_scenario)
+		etdev_warn(etdev, "tpu_performance BTS scenario not found\n");
+	platform_pwr->scenario_count = 0;
+
 	ret = abrolhos_pwr_state_set(etdev,
 				     abrolhos_get_initial_pwr_state(dev));
 	if (ret)
@@ -658,17 +672,24 @@ static int abrolhos_pm_after_create(struct edgetpu_pm *etpm)
 	debugfs_create_file("uart_rate", 0440, abrolhos_pwr_debugfs_dir, dev,
 			    &fops_tpu_uart_rate);
 	debugfs_create_file("policy", 0660, abrolhos_pwr_debugfs_dir,
-			    edgetpu_pdev, &fops_tpu_pwr_policy);
+			    abpdev, &fops_tpu_pwr_policy);
 	debugfs_create_file("core_pwr", 0660, abrolhos_pwr_debugfs_dir,
-			    edgetpu_pdev, &fops_tpu_core_pwr);
+			    abpdev, &fops_tpu_core_pwr);
 
 	return 0;
 }
 
 static void abrolhos_pm_before_destroy(struct edgetpu_pm *etpm)
 {
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
+
 	debugfs_remove_recursive(abrolhos_pwr_debugfs_dir);
 	pm_runtime_disable(etpm->etdev->dev);
+	abrolhos_pm_cleanup_bts_scenario(etdev);
+	exynos_pm_qos_remove_request(&platform_pwr->int_min);
+	exynos_pm_qos_remove_request(&platform_pwr->mif_min);
 }
 
 static struct edgetpu_pm_handlers abrolhos_pm_handlers = {
@@ -680,66 +701,78 @@ static struct edgetpu_pm_handlers abrolhos_pm_handlers = {
 
 int abrolhos_pm_create(struct edgetpu_dev *etdev)
 {
-	exynos_pm_qos_add_request(&int_min, PM_QOS_DEVICE_THROUGHPUT, 0);
-	exynos_pm_qos_add_request(&mif_min, PM_QOS_BUS_THROUGHPUT, 0);
-
-	performance_scenario = bts_get_scenindex("tpu_performance");
-
-	if (!performance_scenario)
-		etdev_warn(etdev, "tpu_performance BTS scenario not found\n");
-
 	return edgetpu_pm_create(etdev, &abrolhos_pm_handlers);
 }
 
 void abrolhos_pm_destroy(struct edgetpu_dev *etdev)
 {
-	abrolhos_pm_cleanup_bts_scenario(etdev);
-	exynos_pm_qos_remove_request(&int_min);
-	exynos_pm_qos_remove_request(&mif_min);
-
 	edgetpu_pm_destroy(etdev);
 }
 
 void abrolhos_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
 {
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
 	s32 int_val = (pm_qos_val >> PM_QOS_INT_SHIFT) * PM_QOS_FACTOR;
 	s32 mif_val = (pm_qos_val & PM_QOS_MIF_MASK) * PM_QOS_FACTOR;
 
 	etdev_dbg(etdev, "%s: pm_qos request - int = %d mif = %d\n", __func__,
 		  int_val, mif_val);
 
-	exynos_pm_qos_update_request(&int_min, int_val);
-	exynos_pm_qos_update_request(&mif_min, mif_val);
+	exynos_pm_qos_update_request(&platform_pwr->int_min, int_val);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, mif_val);
 }
 
 static void abrolhos_pm_activate_bts_scenario(struct edgetpu_dev *etdev)
 {
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
 	/* bts_add_scenario() keeps track of reference count internally.*/
 	int ret;
 
 	if (!performance_scenario)
 		return;
+	mutex_lock(&platform_pwr->scenario_lock);
 	ret = bts_add_scenario(performance_scenario);
 	if (ret)
 		etdev_warn_once(etdev, "error %d adding BTS scenario %u\n", ret,
 				performance_scenario);
 	else
-		atomic64_inc(&scenario_count);
+		platform_pwr->scenario_count++;
+
+	etdev_dbg(etdev, "BTS Scenario activated: %d\n",
+		  platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
 }
 
 static void abrolhos_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 {
 	/* bts_del_scenario() keeps track of reference count internally.*/
 	int ret;
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_pwr *platform_pwr = &abpdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
 
 	if (!performance_scenario)
 		return;
+	mutex_lock(&platform_pwr->scenario_lock);
+	if (!platform_pwr->scenario_count) {
+		etdev_warn(etdev, "Unbalanced bts deactivate\n");
+		mutex_unlock(&platform_pwr->scenario_lock);
+		return;
+	}
 	ret = bts_del_scenario(performance_scenario);
 	if (ret)
 		etdev_warn_once(etdev, "error %d deleting BTS scenario %u\n",
 				ret, performance_scenario);
 	else
-		atomic64_dec(&scenario_count);
+		platform_pwr->scenario_count--;
+
+	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n",
+		  platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
 }
 
 void abrolhos_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val)
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 33ca204..7294dd4 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -37,6 +37,28 @@
 #include "edgetpu-wakelock.h"
 #include "edgetpu.h"
 
+enum edgetpu_vma_type {
+	VMA_INVALID,
+
+	VMA_FULL_CSR,
+	VMA_VII_CSR,
+	VMA_VII_CMDQ,
+	VMA_VII_RESPQ,
+	VMA_LOG,
+	VMA_TRACE,
+};
+
+/* structure to be set to vma->vm_private_data on mmap */
+struct edgetpu_vma_private {
+	struct edgetpu_client *client;
+	enum edgetpu_vma_type type;
+	/*
+	 * vm_private_data is copied when a VMA is split, using this reference
+	 * counter to know when should this object be freed.
+	 */
+	refcount_t count;
+};
+
 static atomic_t single_dev_count = ATOMIC_INIT(-1);
 
 static int edgetpu_mmap_full_csr(struct edgetpu_client *client,
@@ -59,52 +81,133 @@ static int edgetpu_mmap_full_csr(struct edgetpu_client *client,
 	return ret;
 }
 
+static enum edgetpu_vma_type mmap_vma_type(unsigned long pgoff)
+{
+	const unsigned long off = pgoff << PAGE_SHIFT;
+
+	switch (off) {
+	case 0:
+		return VMA_FULL_CSR;
+	case EDGETPU_MMAP_CSR_OFFSET:
+		return VMA_VII_CSR;
+	case EDGETPU_MMAP_CMD_QUEUE_OFFSET:
+		return VMA_VII_CMDQ;
+	case EDGETPU_MMAP_RESP_QUEUE_OFFSET:
+		return VMA_VII_RESPQ;
+	case EDGETPU_MMAP_LOG_BUFFER_OFFSET:
+		return VMA_LOG;
+	case EDGETPU_MMAP_TRACE_BUFFER_OFFSET:
+		return VMA_TRACE;
+	default:
+		return VMA_INVALID;
+	}
+}
+
 /*
- * Returns the wakelock event by mmap offset. Returns EDGETPU_WAKELOCK_EVENT_END
- * if the offset does not correspond to a wakelock event.
+ * Returns the wakelock event by VMA type. Returns EDGETPU_WAKELOCK_EVENT_END
+ * if the type does not correspond to a wakelock event.
  */
-static enum edgetpu_wakelock_event mmap_wakelock_event(unsigned long pgoff)
+static enum edgetpu_wakelock_event
+vma_type_to_wakelock_event(enum edgetpu_vma_type type)
 {
-	switch (pgoff) {
-	case 0:
+	switch (type) {
+	case VMA_FULL_CSR:
 		return EDGETPU_WAKELOCK_EVENT_FULL_CSR;
-	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
+	case VMA_VII_CSR:
 		return EDGETPU_WAKELOCK_EVENT_MBOX_CSR;
-	case EDGETPU_MMAP_CMD_QUEUE_OFFSET >> PAGE_SHIFT:
+	case VMA_VII_CMDQ:
 		return EDGETPU_WAKELOCK_EVENT_CMD_QUEUE;
-	case EDGETPU_MMAP_RESP_QUEUE_OFFSET >> PAGE_SHIFT:
+	case VMA_VII_RESPQ:
 		return EDGETPU_WAKELOCK_EVENT_RESP_QUEUE;
 	default:
 		return EDGETPU_WAKELOCK_EVENT_END;
 	}
 }
 
+static struct edgetpu_vma_private *
+edgetpu_vma_private_alloc(struct edgetpu_client *client,
+			  enum edgetpu_vma_type type)
+{
+	struct edgetpu_vma_private *pvt = kmalloc(sizeof(*pvt), GFP_KERNEL);
+
+	if (!pvt)
+		return NULL;
+	pvt->client = edgetpu_client_get(client);
+	pvt->type = type;
+	refcount_set(&pvt->count, 1);
+
+	return pvt;
+}
+
+static void edgetpu_vma_private_get(struct edgetpu_vma_private *pvt)
+{
+	WARN_ON_ONCE(!refcount_inc_not_zero(&pvt->count));
+}
+
+static void edgetpu_vma_private_put(struct edgetpu_vma_private *pvt)
+{
+	if (!pvt)
+		return;
+	if (refcount_dec_and_test(&pvt->count)) {
+		edgetpu_client_put(pvt->client);
+		kfree(pvt);
+	}
+}
+
 static void edgetpu_vma_open(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_vma_private *pvt = vma->vm_private_data;
+	enum edgetpu_wakelock_event evt;
+	struct edgetpu_client *client;
+	struct edgetpu_dev *etdev;
+
+	edgetpu_vma_private_get(pvt);
+	client = pvt->client;
+	etdev = client->etdev;
 
+	evt = vma_type_to_wakelock_event(pvt->type);
 	if (evt != EDGETPU_WAKELOCK_EVENT_END)
 		edgetpu_wakelock_inc_event(client->wakelock, evt);
+
+	/* handle telemetry types */
+	switch (pvt->type) {
+	case VMA_LOG:
+		edgetpu_telemetry_inc_mmap_count(etdev, EDGETPU_TELEMETRY_LOG);
+		break;
+	case VMA_TRACE:
+		edgetpu_telemetry_inc_mmap_count(etdev,
+						 EDGETPU_TELEMETRY_TRACE);
+		break;
+	default:
+		break;
+	}
 }
 
+/* Records previously mmapped addresses were unmapped. */
 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_vma_private *pvt = vma->vm_private_data;
+	struct edgetpu_client *client = pvt->client;
+	enum edgetpu_wakelock_event evt = vma_type_to_wakelock_event(pvt->type);
 	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);
+	/* handle telemetry types */
+	switch (pvt->type) {
+	case VMA_LOG:
+		edgetpu_telemetry_dec_mmap_count(etdev, EDGETPU_TELEMETRY_LOG);
+		break;
+	case VMA_TRACE:
+		edgetpu_telemetry_dec_mmap_count(etdev,
+						 EDGETPU_TELEMETRY_TRACE);
+		break;
+	default:
+		break;
+	}
+
+	edgetpu_vma_private_put(pvt);
 }
 
 static const struct vm_operations_struct edgetpu_vma_ops = {
@@ -116,7 +219,9 @@ static const struct vm_operations_struct edgetpu_vma_ops = {
 int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 {
 	int ret = 0;
+	enum edgetpu_vma_type type;
 	enum edgetpu_wakelock_event evt;
+	struct edgetpu_vma_private *pvt;
 
 	if (vma->vm_start & ~PAGE_MASK) {
 		etdev_dbg(client->etdev,
@@ -128,61 +233,90 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 	etdev_dbg(client->etdev, "%s: mmap pgoff = %lX\n", __func__,
 		  vma->vm_pgoff);
 
-	vma->vm_private_data = client;
-	vma->vm_ops = &edgetpu_vma_ops;
+	type = mmap_vma_type(vma->vm_pgoff);
+	if (type == VMA_INVALID)
+		return -EINVAL;
+	pvt = edgetpu_vma_private_alloc(client, type);
+	if (!pvt)
+		return -ENOMEM;
 
 	/* Mark the VMA's pages as uncacheable. */
 	vma->vm_page_prot = pgprot_noncached(vma->vm_page_prot);
+	/* Disable fancy things to ensure our event counters work. */
+	vma->vm_flags |= VM_DONTCOPY | VM_DONTEXPAND | VM_DONTDUMP;
 
 	/* map all CSRs for debug purpose */
-	if (!vma->vm_pgoff) {
+	if (type == VMA_FULL_CSR) {
 		evt = EDGETPU_WAKELOCK_EVENT_FULL_CSR;
 		if (edgetpu_wakelock_inc_event(client->wakelock, evt)) {
 			ret = edgetpu_mmap_full_csr(client, vma);
 			if (ret)
-				edgetpu_wakelock_dec_event(client->wakelock,
-							   evt);
+				goto out_dec_evt;
 		} else {
 			ret = -EAGAIN;
 		}
-		return ret;
+		goto out_set_op;
 	}
 
-	/* Allow mapping log and telemetry buffers without creating a group */
-	if (vma->vm_pgoff == EDGETPU_MMAP_LOG_BUFFER_OFFSET >> PAGE_SHIFT)
-		return edgetpu_mmap_telemetry_buffer(
-			client->etdev, EDGETPU_TELEMETRY_LOG, vma);
-	if (vma->vm_pgoff == EDGETPU_MMAP_TRACE_BUFFER_OFFSET >> PAGE_SHIFT)
-		return edgetpu_mmap_telemetry_buffer(
+	/* Allow mapping log and telemetry buffers without a group */
+	if (type == VMA_LOG) {
+		ret = edgetpu_mmap_telemetry_buffer(client->etdev,
+						    EDGETPU_TELEMETRY_LOG, vma);
+		goto out_set_op;
+	}
+	if (type == VMA_TRACE) {
+		ret = edgetpu_mmap_telemetry_buffer(
 			client->etdev, EDGETPU_TELEMETRY_TRACE, vma);
+		goto out_set_op;
+	}
 
-	evt = mmap_wakelock_event(vma->vm_pgoff);
-	if (evt == EDGETPU_WAKELOCK_EVENT_END)
-		return -EINVAL;
-	if (!edgetpu_wakelock_inc_event(client->wakelock, evt))
-		return -EAGAIN;
+	evt = vma_type_to_wakelock_event(type);
+	/*
+	 * @type should always correspond to a valid event since we handled
+	 * telemetry mmaps above, still check evt != END in case new types are
+	 * added in the future.
+	 */
+	if (unlikely(evt == EDGETPU_WAKELOCK_EVENT_END)) {
+		ret = -EINVAL;
+		goto err_release_pvt;
+	}
+	if (!edgetpu_wakelock_inc_event(client->wakelock, evt)) {
+		ret = -EAGAIN;
+		goto err_release_pvt;
+	}
 
 	mutex_lock(&client->group_lock);
 	if (!client->group) {
 		ret = -EINVAL;
 		goto out_unlock;
 	}
-	switch (vma->vm_pgoff) {
-	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
+	switch (type) {
+	case VMA_VII_CSR:
 		ret = edgetpu_mmap_csr(client->group, vma);
 		break;
-	case EDGETPU_MMAP_CMD_QUEUE_OFFSET >> PAGE_SHIFT:
+	case VMA_VII_CMDQ:
 		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma);
 		break;
-	case EDGETPU_MMAP_RESP_QUEUE_OFFSET >> PAGE_SHIFT:
+	case VMA_VII_RESPQ:
 		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE,
 					 vma);
 		break;
+	default: /* to appease compiler */
+		break;
 	}
 out_unlock:
 	mutex_unlock(&client->group_lock);
+out_dec_evt:
 	if (ret)
 		edgetpu_wakelock_dec_event(client->wakelock, evt);
+out_set_op:
+	if (!ret) {
+		vma->vm_private_data = pvt;
+		vma->vm_ops = &edgetpu_vma_ops;
+		return 0;
+	}
+err_release_pvt:
+	edgetpu_vma_private_put(pvt);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 5e2a9d0..ba13c85 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -1157,8 +1157,8 @@ error:
  */
 static struct edgetpu_host_map *
 alloc_mapping_from_useraddr(struct edgetpu_device_group *group, u64 host_addr,
-			    u64 size, edgetpu_map_flag_t flags,
-			    struct page **pages, uint num_pages)
+			    edgetpu_map_flag_t flags, struct page **pages,
+			    uint num_pages)
 {
 	struct edgetpu_dev *etdev = group->etdev;
 	struct edgetpu_host_map *hmap;
@@ -1199,9 +1199,9 @@ alloc_mapping_from_useraddr(struct edgetpu_device_group *group, u64 host_addr,
 			sgt = &hmap->map.sgt;
 		else
 			sgt = &hmap->sg_tables[i];
-		ret = sg_alloc_table_from_pages(sgt, pages, num_pages,
-						host_addr & (PAGE_SIZE - 1),
-						size, GFP_KERNEL);
+		ret = sg_alloc_table_from_pages(sgt, pages, num_pages, 0,
+						num_pages * PAGE_SIZE,
+						GFP_KERNEL);
 		if (ret) {
 			etdev_dbg(etdev,
 				  "%s: sg_alloc_table_from_pages failed %u:%pK-%u: %d",
@@ -1223,8 +1223,6 @@ error_free_sgt:
 		sg_free_table(sgt);
 	}
 error:
-	for (i = 0; i < num_pages; i++)
-		unpin_user_page(pages[i]);
 	if (hmap) {
 		edgetpu_device_group_put(hmap->map.priv);
 		kfree(hmap->sg_tables);
@@ -1306,14 +1304,16 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 	struct page **pages;
 	int ret = -EINVAL;
 	u64 host_addr = arg->host_address;
-	u64 size = arg->size;
 	edgetpu_map_flag_t flags = arg->flags;
 	struct edgetpu_host_map *hmap;
-	struct edgetpu_mapping *map;
+	struct edgetpu_mapping *map = NULL;
 	struct edgetpu_dev *etdev;
 	enum edgetpu_context_id context_id;
 	const u32 mmu_flags = map_to_mmu_flags(flags) | EDGETPU_MMU_HOST;
+	int i;
 
+	if (!valid_dma_direction(flags & EDGETPU_MAP_DIR_MASK))
+		return -EINVAL;
 	/* Pin user pages before holding any lock. */
 	pages = edgetpu_pin_user_pages(group, arg, &num_pages);
 	if (IS_ERR(pages))
@@ -1323,20 +1323,20 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 	context_id = edgetpu_group_context_id_locked(group);
 	if (!edgetpu_device_group_is_finalized(group)) {
 		ret = -EINVAL;
-		goto error_unlock_group;
+		goto error;
 	}
 	if (!IS_MIRRORED(flags)) {
 		if (arg->die_index >= group->n_clients) {
 			ret = -EINVAL;
-			goto error_unlock_group;
+			goto error;
 		}
 	}
 
-	hmap = alloc_mapping_from_useraddr(group, host_addr, size, flags, pages,
+	hmap = alloc_mapping_from_useraddr(group, host_addr, flags, pages,
 					   num_pages);
 	if (IS_ERR(hmap)) {
 		ret = PTR_ERR(hmap);
-		goto error_unlock_group;
+		goto error;
 	}
 
 	map = &hmap->map;
@@ -1345,27 +1345,27 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 		etdev = group->etdev;
 		ret = edgetpu_mmu_map(etdev, map, context_id, mmu_flags);
 		if (ret)
-			goto error_release_map;
+			goto error;
 		ret = edgetpu_device_group_map_iova_sgt(group, hmap);
 		if (ret) {
 			etdev_dbg(etdev,
 				  "group add translation failed %u:0x%llx",
 				  group->workload_id, map->device_address);
-			goto error_release_map;
+			goto error;
 		}
 	} else {
 		map->die_index = arg->die_index;
 		etdev = edgetpu_device_group_nth_etdev(group, map->die_index);
 		ret = edgetpu_mmu_map(etdev, map, context_id, mmu_flags);
 		if (ret)
-			goto error_release_map;
+			goto error;
 	}
 
 	ret = edgetpu_mapping_add(&group->host_mappings, map);
 	if (ret) {
 		etdev_dbg(etdev, "duplicate mapping %u:0x%llx",
 			  group->workload_id, map->device_address);
-		goto error_release_map;
+		goto error;
 	}
 
 	mutex_unlock(&group->lock);
@@ -1373,13 +1373,17 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 	kfree(pages);
 	return 0;
 
-error_release_map:
-	edgetpu_mapping_lock(&group->host_mappings);
-	/* this will free @hmap */
-	edgetpu_unmap_node(map);
-	edgetpu_mapping_unlock(&group->host_mappings);
-
-error_unlock_group:
+error:
+	if (map) {
+		edgetpu_mapping_lock(&group->host_mappings);
+		/* this will free @hmap */
+		edgetpu_unmap_node(map);
+		edgetpu_mapping_unlock(&group->host_mappings);
+	} else {
+		/* revert edgetpu_pin_user_pages() */
+		for (i = 0; i < num_pages; i++)
+			unpin_user_page(pages[i]);
+	}
 	mutex_unlock(&group->lock);
 	kfree(pages);
 	return ret;
@@ -1431,6 +1435,8 @@ int edgetpu_device_group_sync_buffer(struct edgetpu_device_group *group,
 	enum dma_data_direction dir = arg->flags & EDGETPU_MAP_DIR_MASK;
 	struct edgetpu_host_map *hmap;
 
+	if (!valid_dma_direction(dir))
+		return -EINVAL;
 	/* invalid if size == 0 or overflow */
 	if (arg->offset + arg->size <= arg->offset)
 		return -EINVAL;
@@ -1547,11 +1553,7 @@ int edgetpu_mmap_queue(struct edgetpu_device_group *group,
 	edgetpu_queue_mem *queue_mem;
 
 	mutex_lock(&group->lock);
-	/*
-	 * VII queues are available even when mailbox detached, no need to check
-	 * whether mailbox attached here.
-	 */
-	if (!edgetpu_device_group_is_finalized(group)) {
+	if (!edgetpu_group_finalized_and_attached(group)) {
 		ret = -EINVAL;
 		goto out;
 	}
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index b0e3886..f5c5a62 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -7,6 +7,7 @@
 
 #include <linux/debugfs.h>
 #include <linux/dma-buf.h>
+#include <linux/dma-direction.h>
 #include <linux/dma-fence.h>
 #include <linux/dma-mapping.h>
 #include <linux/ktime.h>
@@ -670,8 +671,8 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	tpu_addr_t tpu_addr;
 	uint i;
 
-	/* offset is not page-aligned */
-	if (offset_in_page(offset))
+	/* invalid DMA direction or offset is not page-aligned */
+	if (!valid_dma_direction(dir) || offset_in_page(offset))
 		return -EINVAL;
 	/* size == 0 or overflow */
 	if (offset + size <= offset)
@@ -786,7 +787,7 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 	tpu_addr_t tpu_addr;
 	int i;
 
-	if (arg->size == 0)
+	if (!valid_dma_direction(dir) || arg->size == 0)
 		return -EINVAL;
 	mutex_lock(&group->lock);
 	if (!edgetpu_device_group_is_finalized(group))
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 673a51c..ded7cc7 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -472,11 +472,17 @@ int edgetpu_firmware_restart_locked(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_firmware *et_fw = etdev->firmware;
 	const struct edgetpu_firmware_handlers *handlers = et_fw->p->handlers;
-	int ret;
+	int ret = -1;
 
 	et_fw->p->status = FW_LOADING;
 	edgetpu_sw_wdt_stop(etdev);
-	if (handlers && handlers->prepare_run) {
+	/*
+	 * Try restarting the firmware first, fall back to normal firmware start
+	 * if this fails.
+	 */
+	if (handlers && handlers->restart)
+		ret = handlers->restart(et_fw);
+	if (ret && handlers && handlers->prepare_run) {
 		ret = handlers->prepare_run(et_fw, &et_fw->p->fw_desc.buf);
 		if (ret)
 			return ret;
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index 2303311..477d1a5 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -147,6 +147,12 @@ struct edgetpu_firmware_handlers {
 			   struct edgetpu_firmware_buffer *fw_buf);
 	/* Firmware running, after successful handshake. */
 	void (*launch_complete)(struct edgetpu_firmware *et_fw);
+
+	/*
+	 * Optional platform-specific handler to restart an already loaded
+	 * firmware.
+	 */
+	int (*restart)(struct edgetpu_firmware *et_fw);
 };
 
 /*
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index b62ecea..f48fe20 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -235,7 +235,7 @@ out:
 int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 {
 #ifdef CONFIG_ABROLHOS
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 #endif
 	struct edgetpu_iommu *etiommu;
 	int ret;
@@ -272,17 +272,15 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	etdev->mmu_cookie = etiommu;
 	/* TODO (b/178571278): remove chipset specific code. */
 #ifdef CONFIG_ABROLHOS
-	if (!edgetpu_pdev->csr_iova)
+	if (!abpdev->csr_iova)
 		goto success;
 
 	etdev_dbg(etdev, "Mapping device CSRs: %llX -> %llX (%lu bytes)\n",
-		  edgetpu_pdev->csr_iova, edgetpu_pdev->csr_paddr,
-		  edgetpu_pdev->csr_size);
+		  abpdev->csr_iova, abpdev->csr_paddr, abpdev->csr_size);
 
 	/* Add an IOMMU translation for the CSR region */
-	ret = edgetpu_mmu_add_translation(etdev, edgetpu_pdev->csr_iova,
-					  edgetpu_pdev->csr_paddr,
-					  edgetpu_pdev->csr_size,
+	ret = edgetpu_mmu_add_translation(etdev, abpdev->csr_iova,
+					  abpdev->csr_paddr, abpdev->csr_size,
 					  IOMMU_READ | IOMMU_WRITE | IOMMU_PRIV,
 					  EDGETPU_CONTEXT_KCI);
 	if (ret) {
@@ -309,7 +307,7 @@ void edgetpu_mmu_reset(struct edgetpu_dev *etdev)
 void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
 {
 #ifdef CONFIG_ABROLHOS
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
 #endif
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	int i, ret;
@@ -318,13 +316,13 @@ void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
 		return;
 
 #ifdef CONFIG_ABROLHOS
-	if (edgetpu_pdev->csr_iova) {
-		edgetpu_mmu_remove_translation(&edgetpu_pdev->edgetpu_dev,
-					       edgetpu_pdev->csr_iova,
-					       edgetpu_pdev->csr_size,
+	if (abpdev->csr_iova) {
+		edgetpu_mmu_remove_translation(&abpdev->edgetpu_dev,
+					       abpdev->csr_iova,
+					       abpdev->csr_size,
 					       EDGETPU_CONTEXT_KCI);
 	}
-	edgetpu_pdev->csr_iova = 0;
+	abpdev->csr_iova = 0;
 #endif
 	ret = edgetpu_unregister_iommu_device_fault_handler(etdev);
 	if (ret)
diff --git a/drivers/edgetpu/edgetpu-iremap-pool.c b/drivers/edgetpu/edgetpu-iremap-pool.c
index b23bcc6..91b8fea 100644
--- a/drivers/edgetpu/edgetpu-iremap-pool.c
+++ b/drivers/edgetpu/edgetpu-iremap-pool.c
@@ -133,6 +133,7 @@ int edgetpu_iremap_mmap(struct edgetpu_dev *etdev, struct vm_area_struct *vma,
 	phys_addr_t phys;
 	int ret;
 	unsigned long orig_pgoff = vma->vm_pgoff;
+	ulong vma_size, map_size;
 
 #ifdef CONFIG_ARM64
 	/*
@@ -157,8 +158,10 @@ int edgetpu_iremap_mmap(struct edgetpu_dev *etdev, struct vm_area_struct *vma,
 	phys = etmempool->base_phys_addr + offset;
 	etdev_dbg(etdev, "%s: virt = %llx phys = %llx\n",
 		  __func__, (u64)mem->vaddr, phys);
+	vma_size = vma->vm_end - vma->vm_start;
+	map_size = min(vma_size, mem->size);
 	ret = remap_pfn_range(vma, vma->vm_start, phys >> PAGE_SHIFT,
-			      vma->vm_end - vma->vm_start, vma->vm_page_prot);
+			      map_size, vma->vm_page_prot);
 	vma->vm_pgoff = orig_pgoff;
 	return ret;
 }
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index 0ade5b0..cf9435c 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -250,6 +250,16 @@ static void telemetry_mappings_show(struct edgetpu_telemetry *tel,
 		   tel->coherent_mem.host_addr, &tel->coherent_mem.dma_addr);
 }
 
+static void telemetry_inc_mmap_count(struct edgetpu_telemetry *tel, int dif)
+{
+	if (!tel->inited)
+		return;
+
+	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)
@@ -261,16 +271,17 @@ static int telemetry_mmap_buffer(struct edgetpu_dev *etdev,
 
 	mutex_lock(&tel->mmap_lock);
 
-	if (!tel->is_mmapped) {
+	if (!tel->mmapped_count) {
 		ret = edgetpu_iremap_mmap(etdev, vma, &tel->coherent_mem);
 
 		if (!ret) {
 			tel->coherent_mem.host_addr = vma->vm_start;
-			tel->is_mmapped = true;
+			tel->mmapped_count = 1;
 		}
 	} else {
 		ret = -EBUSY;
-		etdev_warn(etdev, "Buffer is already mmapped");
+		etdev_warn(etdev, "%s is already mmapped %ld times", tel->name,
+			   tel->mmapped_count);
 	}
 
 	mutex_unlock(&tel->mmap_lock);
@@ -278,18 +289,6 @@ static int telemetry_mmap_buffer(struct edgetpu_dev *etdev,
 	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,
@@ -344,7 +343,7 @@ static int telemetry_init(struct edgetpu_dev *etdev,
 	tel->state = EDGETPU_TELEMETRY_ENABLED;
 	tel->inited = true;
 	mutex_init(&tel->mmap_lock);
-	tel->is_mmapped = false;
+	tel->mmapped_count = 0;
 
 	return 0;
 }
@@ -467,12 +466,18 @@ int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev,
 		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)
+void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev,
+				      enum edgetpu_telemetry_type type)
+{
+	if (!etdev->telemetry)
+		return;
+	telemetry_inc_mmap_count(select_telemetry(etdev->telemetry, type), 1);
+}
+
+void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev,
+				      enum edgetpu_telemetry_type type)
 {
 	if (!etdev->telemetry)
 		return;
-	telemetry_munmap_buffer(etdev, select_telemetry(etdev->telemetry, type),
-				vma);
+	telemetry_inc_mmap_count(select_telemetry(etdev->telemetry, type), -1);
 }
diff --git a/drivers/edgetpu/edgetpu-telemetry.h b/drivers/edgetpu/edgetpu-telemetry.h
index 0a57601..a25d002 100644
--- a/drivers/edgetpu/edgetpu-telemetry.h
+++ b/drivers/edgetpu/edgetpu-telemetry.h
@@ -86,9 +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 mutex mmap_lock; /* protects mmapped_count */
+	/* number of VMAs that are mapped to this telemetry buffer */
+	long mmapped_count;
 };
 
 struct edgetpu_telemetry_ctx {
@@ -142,8 +142,9 @@ 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);
+void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev,
+				      enum edgetpu_telemetry_type type);
+void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev,
+				      enum edgetpu_telemetry_type type);
 
 #endif /* __EDGETPU_TELEMETRY_H__ */
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index 421cd29..1404674 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -147,7 +147,54 @@ static void edgetpu_counter_update(
 
 	mutex_lock(&ustats->usage_stats_lock);
 	if (counter->type >= 0 && counter->type < EDGETPU_COUNTER_COUNT)
-		ustats->counter[counter->type] = counter->value;
+		ustats->counter[counter->type] += counter->value;
+	mutex_unlock(&ustats->usage_stats_lock);
+}
+
+static void edgetpu_max_watermark_update(
+	struct edgetpu_dev *etdev,
+	struct edgetpu_usage_max_watermark *max_watermark)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (!ustats)
+		return;
+
+	etdev_dbg(etdev, "%s: type=%d value=%llu\n", __func__,
+		  max_watermark->type, max_watermark->value);
+
+	if (max_watermark->type < 0 ||
+	    max_watermark->type >= EDGETPU_MAX_WATERMARK_TYPE_COUNT)
+		return;
+
+	mutex_lock(&ustats->usage_stats_lock);
+	if (max_watermark->value > ustats->max_watermark[max_watermark->type])
+		ustats->max_watermark[max_watermark->type] =
+			max_watermark->value;
+	mutex_unlock(&ustats->usage_stats_lock);
+}
+
+static void edgetpu_thread_stats_update(
+	struct edgetpu_dev *etdev,
+	struct edgetpu_thread_stats *thread_stats)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (!ustats)
+		return;
+
+	etdev_dbg(etdev, "%s: id=%d stackmax=%u\n", __func__,
+		  thread_stats->thread_id, thread_stats->max_stack_usage_bytes);
+
+	if (thread_stats->thread_id < 0 ||
+	    thread_stats->thread_id >= EDGETPU_FW_THREAD_COUNT)
+		return;
+
+	mutex_lock(&ustats->usage_stats_lock);
+	if (thread_stats->max_stack_usage_bytes >
+	    ustats->thread_stack_max[thread_stats->thread_id])
+		ustats->thread_stack_max[thread_stats->thread_id] =
+			thread_stats->max_stack_usage_bytes;
 	mutex_unlock(&ustats->usage_stats_lock);
 }
 
@@ -178,6 +225,14 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 		case EDGETPU_METRIC_TYPE_COUNTER:
 			edgetpu_counter_update(etdev, &metric->counter);
 			break;
+		case EDGETPU_METRIC_TYPE_MAX_WATERMARK:
+			edgetpu_max_watermark_update(
+				etdev, &metric->max_watermark);
+			break;
+		case EDGETPU_METRIC_TYPE_THREAD_STATS:
+			edgetpu_thread_stats_update(
+				etdev, &metric->thread_stats);
+			break;
 		default:
 			etdev_dbg(etdev, "%s: %d: skip unknown type=%u",
 				  __func__, i, metric->type);
@@ -209,7 +264,7 @@ static int64_t edgetpu_usage_get_counter(
 	enum edgetpu_usage_counter_type counter_type)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
-	int32_t val;
+	int64_t val;
 
 	if (counter_type >= EDGETPU_COUNTER_COUNT)
 		return -1;
@@ -220,6 +275,22 @@ static int64_t edgetpu_usage_get_counter(
 	return val;
 }
 
+static int64_t edgetpu_usage_get_max_watermark(
+	struct edgetpu_dev *etdev,
+	enum edgetpu_usage_max_watermark_type max_watermark_type)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int64_t val;
+
+	if (max_watermark_type >= EDGETPU_MAX_WATERMARK_TYPE_COUNT)
+		return -1;
+	edgetpu_kci_update_usage(etdev);
+	mutex_lock(&ustats->usage_stats_lock);
+	val = ustats->max_watermark[max_watermark_type];
+	mutex_unlock(&ustats->usage_stats_lock);
+	return val;
+}
+
 static ssize_t tpu_usage_show(struct device *dev,
 			      struct device_attribute *attr,
 			      char *buf)
@@ -347,12 +418,163 @@ static ssize_t tpu_throttle_stall_count_show(struct device *dev,
 }
 static DEVICE_ATTR_RO(tpu_throttle_stall_count);
 
+static ssize_t inference_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_INFERENCES);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(inference_count);
+
+static ssize_t tpu_op_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_OPS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(tpu_op_count);
+
+static ssize_t param_cache_hit_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_PARAM_CACHE_HITS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(param_cache_hit_count);
+
+static ssize_t param_cache_miss_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_PARAM_CACHE_MISSES);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(param_cache_miss_count);
+
+static ssize_t context_preempt_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_CONTEXT_PREEMPTS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+static DEVICE_ATTR_RO(context_preempt_count);
+
+static ssize_t outstanding_commands_max_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_max_watermark(
+			etdev, EDGETPU_MAX_WATERMARK_OUT_CMDS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+
+static ssize_t outstanding_commands_max_store(
+	struct device *dev, struct device_attribute *attr,
+	const char *buf, size_t count)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (ustats) {
+		mutex_lock(&ustats->usage_stats_lock);
+		ustats->max_watermark[EDGETPU_MAX_WATERMARK_OUT_CMDS] = 0;
+		mutex_unlock(&ustats->usage_stats_lock);
+	}
+
+	return count;
+}
+static DEVICE_ATTR_RW(outstanding_commands_max);
+
+static ssize_t preempt_depth_max_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_max_watermark(
+			etdev, EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+
+static ssize_t preempt_depth_max_store(
+	struct device *dev, struct device_attribute *attr,
+	const char *buf, size_t count)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (ustats) {
+		mutex_lock(&ustats->usage_stats_lock);
+		ustats->max_watermark[EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH] = 0;
+		mutex_unlock(&ustats->usage_stats_lock);
+	}
+
+	return count;
+}
+static DEVICE_ATTR_RW(preempt_depth_max);
+
+static ssize_t fw_thread_stats_show(
+	struct device *dev, struct device_attribute *attr, char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int i;
+	ssize_t ret = 0;
+
+	edgetpu_kci_update_usage(etdev);
+	mutex_lock(&ustats->usage_stats_lock);
+
+	for (i = 0; i < EDGETPU_FW_THREAD_COUNT; i++) {
+		if (!ustats->thread_stack_max[i])
+			continue;
+		ret += scnprintf(buf + ret, PAGE_SIZE - ret,
+				 "%u\t%u\n", i, ustats->thread_stack_max[i]);
+		/* Not checking ret < PAGE_SIZE is intended. */
+	}
+
+	mutex_unlock(&ustats->usage_stats_lock);
+	return ret;
+}
+static DEVICE_ATTR_RO(fw_thread_stats);
+
 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,
+	&dev_attr_inference_count.attr,
+	&dev_attr_tpu_op_count.attr,
+	&dev_attr_param_cache_hit_count.attr,
+	&dev_attr_param_cache_miss_count.attr,
+	&dev_attr_context_preempt_count.attr,
+	&dev_attr_outstanding_commands_max.attr,
+	&dev_attr_preempt_depth_max.attr,
+	&dev_attr_fw_thread_stats.attr,
 	NULL,
 };
 
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index a971ad7..7ea3e9d 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -62,8 +62,18 @@ enum edgetpu_usage_counter_type {
 	EDGETPU_COUNTER_TPU_ACTIVE_CYCLES = 0,
 	/* Number of stalls caused by throttling. */
 	EDGETPU_COUNTER_TPU_THROTTLE_STALLS = 1,
+	/* Number of graph invocations. */
+	EDGETPU_COUNTER_INFERENCES = 2,
+	/* Number of TPU offload op invocations. */
+	EDGETPU_COUNTER_TPU_OPS = 3,
+	/* Number of times a TPU op invocation used its cached parameters. */
+	EDGETPU_COUNTER_PARAM_CACHE_HITS = 4,
+	/* Number of times a TPU op invocation had to cache its parameters. */
+	EDGETPU_COUNTER_PARAM_CACHE_MISSES = 5,
+	/* Number of times a context got preempted by another. */
+	EDGETPU_COUNTER_CONTEXT_PREEMPTS = 6,
 
-	EDGETPU_COUNTER_COUNT = 2, /* number of counters above */
+	EDGETPU_COUNTER_COUNT = 7, /* number of counters above */
 };
 
 /* Generic counter. Only reported if it has a value larger than 0. */
@@ -75,12 +85,68 @@ struct __packed edgetpu_usage_counter {
 	uint64_t value;
 };
 
+/* Defines different max watermarks we track. */
+/* Must be kept in sync with firmware MaxWatermarkType */
+enum edgetpu_usage_max_watermark_type {
+	/* Number of outstanding commands in VII trackers of all contexts. */
+	EDGETPU_MAX_WATERMARK_OUT_CMDS = 0,
+	/* Number of preempted contexts at any given time. */
+	EDGETPU_MAX_WATERMARK_PREEMPT_DEPTH = 1,
+
+	/* Number of watermark types above */
+	EDGETPU_MAX_WATERMARK_TYPE_COUNT = 2,
+};
+
+/* Max watermark. Only reported if it has a value larger than 0. */
+struct __packed edgetpu_usage_max_watermark {
+	/* What it counts. */
+	enum edgetpu_usage_max_watermark_type type;
+
+	/*
+	 * Maximum value since last initialization (virtual device join in
+	 * non-mobile, firmware boot on mobile).
+	 */
+	uint64_t value;
+};
+
+/* An enum to identify the tracked firmware threads. */
+/* Must be kept in sync with firmware enum class UsageTrackerThreadId. */
+enum edgetpu_usage_threadid {
+	EDGETPU_FW_THREAD_MAIN = 0,
+	EDGETPU_FW_THREAD_KCI_HANDLER = 1,
+	EDGETPU_FW_THREAD_POWER_ADMIN = 2,
+	EDGETPU_FW_THREAD_VII_SCHEDULER = 3,
+	EDGETPU_FW_THREAD_VII_HANDLER = 4,
+	EDGETPU_FW_THREAD_MCP_GRAPH_DRIVER = 5,
+	EDGETPU_FW_THREAD_SCP_GRAPH_DRIVER = 6,
+	EDGETPU_FW_THREAD_TPU_DRIVER = 7,
+	EDGETPU_FW_THREAD_RESTART_HANDLER = 8,
+	EDGETPU_FW_THREAD_POLL_SERVICE = 9,
+	EDGETPU_FW_THREAD_DMA_DRIVER = 10,
+	EDGETPU_FW_THREAD_GRAPH_DMA_DRIVER = 11,
+
+	/* Number of task identifiers above. */
+	EDGETPU_FW_THREAD_COUNT = 12,
+};
+
+/* Statistics related to a single thread in firmware. */
+/* Must be kept in sync with firmware struct ThreadStats. */
+struct edgetpu_thread_stats {
+	/* The thread in question. */
+	enum edgetpu_usage_threadid thread_id;
+
+	/* Maximum stack usage (in bytes) since last firmware boot. */
+	uint32_t max_stack_usage_bytes;
+};
+
 /* 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,
+	EDGETPU_METRIC_TYPE_THREAD_STATS = 4,
+	EDGETPU_METRIC_TYPE_MAX_WATERMARK = 5,
 };
 
 /*
@@ -94,6 +160,8 @@ struct edgetpu_usage_metric {
 		struct tpu_usage tpu_usage;
 		struct edgetpu_component_activity component_activity;
 		struct edgetpu_usage_counter counter;
+		struct edgetpu_thread_stats thread_stats;
+		struct edgetpu_usage_max_watermark max_watermark;
 	};
 };
 
@@ -104,6 +172,8 @@ struct edgetpu_usage_stats {
 	/* component utilization values reported by firmware */
 	int32_t component_utilization[EDGETPU_USAGE_COMPONENT_COUNT];
 	int64_t counter[EDGETPU_COUNTER_COUNT];
+	int64_t max_watermark[EDGETPU_MAX_WATERMARK_TYPE_COUNT];
+	int32_t thread_stack_max[EDGETPU_FW_THREAD_COUNT];
 	struct mutex usage_stats_lock;
 };