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

* whitechapel:
  edgetpu: set cacheability according to device
  edgetpu: add offset back to dmabuf mapping
  edgetpu: refactor error handling
  edgetpu: reset thread stats on write
  edgetpu: abrolhos remove parsing csr-iova nodes
  edgetpu: remove mapping CSRs on mmu setup
  edgetpu: add corrupted log queue msg back
  edgetpu: deprecate offset in edgetpu_map_dmabuf_ioctl
  edgetpu: fix use-after-free Read in edgetpu_chip_acquire_ext_mailbox
  edgetpu: map dmabuf ignore size argument

Bug: 189968759

Change-Id: I1a0c652eeb002e9f43639dec945a43dadf7b39d2

14 files changed, 121 insertions, 173 deletions

diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index 1f7bdc3..a867c13 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -240,3 +240,13 @@ int edgetpu_chip_release_ext_mailbox(struct edgetpu_client *client,
 	mutex_unlock(&apdev->tz_mailbox_lock);
 	return ret;
 }
+
+void edgetpu_chip_client_remove(struct edgetpu_client *client)
+{
+	struct abrolhos_platform_dev *apdev = to_abrolhos_dev(client->etdev);
+
+	mutex_lock(&apdev->tz_mailbox_lock);
+	if (apdev->secure_client == client)
+		apdev->secure_client = NULL;
+	mutex_unlock(&apdev->tz_mailbox_lock);
+}
diff --git a/drivers/edgetpu/abrolhos-platform.c b/drivers/edgetpu/abrolhos-platform.c
index 16e36dc..caf0910 100644
--- a/drivers/edgetpu/abrolhos-platform.c
+++ b/drivers/edgetpu/abrolhos-platform.c
@@ -82,7 +82,6 @@ edgetpu_platform_setup_fw_region(struct abrolhos_platform_dev *etpdev)
 	struct resource r;
 	struct device_node *np;
 	int err;
-	u32 csr_phys, csr_iova, csr_size;
 	size_t region_map_size =
 		EDGETPU_FW_SIZE_MAX + EDGETPU_REMAPPED_DATA_SIZE;
 
@@ -133,32 +132,7 @@ edgetpu_platform_setup_fw_region(struct abrolhos_platform_dev *etpdev)
 	}
 	etpdev->shared_mem_paddr = r.start + EDGETPU_REMAPPED_DATA_OFFSET;
 
-	err = of_property_read_u32(dev->of_node, "csr-iova", &csr_iova);
-	/* Device did not define a CSR region */
-	if (err)
-		return 0;
-
-	/* If an IOVA was found, we must also have physical address and size */
-	err = of_property_read_u32(dev->of_node, "csr-phys", &csr_phys);
-	if (err) {
-		dev_err(dev, "Device tree: invalid CSR physical address\n");
-		goto out_unmap;
-	}
-
-	err = of_property_read_u32(dev->of_node, "csr-size", &csr_size);
-	if (err) {
-		dev_err(dev, "Device tree: invalid CSR size\n");
-		goto out_unmap;
-	}
-
-	etpdev->csr_paddr = csr_phys;
-	etpdev->csr_iova = csr_iova;
-	etpdev->csr_size = csr_size;
 	return 0;
-out_unmap:
-	memunmap(etpdev->shared_mem_vaddr);
-	etpdev->shared_mem_vaddr = NULL;
-	return err;
 }
 
 static void edgetpu_platform_cleanup_fw_region(
@@ -172,7 +146,7 @@ static void edgetpu_platform_cleanup_fw_region(
 	etpdev->shared_mem_vaddr = NULL;
 }
 
-int edgetpu_setup_mmu(struct edgetpu_dev *etdev)
+int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev)
 {
 	int ret;
 
@@ -183,6 +157,11 @@ int edgetpu_setup_mmu(struct edgetpu_dev *etdev)
 	return ret;
 }
 
+void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
+{
+	edgetpu_mmu_detach(etdev);
+}
+
 static int abrolhos_parse_ssmt(struct abrolhos_platform_dev *etpdev)
 {
 	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
diff --git a/drivers/edgetpu/abrolhos-platform.h b/drivers/edgetpu/abrolhos-platform.h
index 039a44d..73e2221 100644
--- a/drivers/edgetpu/abrolhos-platform.h
+++ b/drivers/edgetpu/abrolhos-platform.h
@@ -47,9 +47,6 @@ struct abrolhos_platform_dev {
 	void *shared_mem_vaddr;
 	phys_addr_t shared_mem_paddr;
 	size_t shared_mem_size;
-	phys_addr_t csr_paddr;
-	dma_addr_t csr_iova;
-	size_t csr_size;
 	struct device *gsa_dev;
 	void __iomem *ssmt_base;
 	struct edgetpu_coherent_mem log_mem;
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index f1158a0..0820e95 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -389,7 +389,7 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 			etdev->dev_name, ret);
 		goto remove_dev;
 	}
-	ret = edgetpu_setup_mmu(etdev);
+	ret = edgetpu_chip_setup_mmu(etdev);
 	if (ret)
 		goto remove_dev;
 
@@ -398,20 +398,20 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 	etdev->kci = devm_kzalloc(etdev->dev, sizeof(*etdev->kci), GFP_KERNEL);
 	if (!etdev->kci) {
 		ret = -ENOMEM;
-		goto detach_mmu;
+		goto remove_usage_stats;
 	}
 
 	etdev->telemetry =
 		devm_kzalloc(etdev->dev, sizeof(*etdev->telemetry), GFP_KERNEL);
 	if (!etdev->telemetry) {
 		ret = -ENOMEM;
-		goto detach_mmu;
+		goto remove_usage_stats;
 	}
 
 	ret = edgetpu_kci_init(etdev->mailbox_manager, etdev->kci);
 	if (ret) {
 		etdev_err(etdev, "edgetpu_kci_init returns %d\n", ret);
-		goto detach_mmu;
+		goto remove_usage_stats;
 	}
 
 	ret = edgetpu_device_dram_init(etdev);
@@ -432,9 +432,9 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 remove_kci:
 	/* releases the resources of KCI */
 	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
-detach_mmu:
+remove_usage_stats:
 	edgetpu_usage_stats_exit(etdev);
-	edgetpu_mmu_detach(etdev);
+	edgetpu_chip_remove_mmu(etdev);
 remove_dev:
 	edgetpu_mark_probe_fail(etdev);
 	edgetpu_fs_remove(etdev);
@@ -448,7 +448,7 @@ void edgetpu_device_remove(struct edgetpu_dev *etdev)
 	edgetpu_device_dram_exit(etdev);
 	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
 	edgetpu_usage_stats_exit(etdev);
-	edgetpu_mmu_detach(etdev);
+	edgetpu_chip_remove_mmu(etdev);
 	edgetpu_fs_remove(etdev);
 }
 
@@ -522,6 +522,7 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	    1 << perdie_event_id_to_num(EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE))
 		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_TRACE);
 
+	edgetpu_chip_client_remove(client);
 	edgetpu_client_put(client);
 }
 
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 987e4ab..3b2e6fc 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -469,7 +469,6 @@ static bool edgetpu_clients_groupable(const struct edgetpu_client *client1,
 {
 	struct edgetpu_dev *etdev1 = client1->etdev, *etdev2 = client2->etdev;
 
-	/* TODO(b/159394046): perform more checks */
 	return etdev1->mcp_id == etdev2->mcp_id &&
 	       etdev1->mcp_die_index != etdev2->mcp_die_index;
 }
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 03918ff..a540763 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -97,7 +97,7 @@ static int etdev_add_translations(struct edgetpu_dev *etdev,
 				  enum dma_data_direction dir,
 				  enum edgetpu_context_id ctx_id)
 {
-	const int prot = __dma_dir_to_iommu_prot(dir);
+	const int prot = __dma_dir_to_iommu_prot(dir, etdev->dev);
 	uint i;
 	u64 offset = 0;
 	int ret;
@@ -280,11 +280,12 @@ static void dmabuf_map_callback_release(struct edgetpu_mapping *map)
 		container_of(map, struct edgetpu_dmabuf_map, map);
 	struct edgetpu_device_group *group = map->priv;
 	const enum dma_data_direction dir = edgetpu_host_dma_dir(map->dir);
-	const tpu_addr_t tpu_addr = map->device_address;
+	tpu_addr_t tpu_addr = map->device_address;
 	struct edgetpu_dev *etdev;
 	uint i;
 
 	if (tpu_addr) {
+		tpu_addr -= dmap->offset;
 		if (IS_MIRRORED(map->flags)) {
 			group_unmap_dmabuf(group, dmap, tpu_addr);
 		} else {
@@ -339,19 +340,13 @@ static void dmabuf_map_callback_show(struct edgetpu_mapping *map,
 		container_of(map, struct edgetpu_dmabuf_map, map);
 
 	if (IS_MIRRORED(dmap->map.flags))
-		seq_printf(
-			s,
-			"  <%s> mirrored: iova=0x%llx pages=%llu %s offset=0x%llx",
-			dmap->dmabufs[0]->exp_name, map->device_address,
-			dmap->size / PAGE_SIZE, edgetpu_dma_dir_rw_s(map->dir),
-			dmap->offset);
+		seq_printf(s, "  <%s> mirrored: iova=0x%llx pages=%llu %s",
+			   dmap->dmabufs[0]->exp_name, map->device_address, dmap->size / PAGE_SIZE,
+			   edgetpu_dma_dir_rw_s(map->dir));
 	else
-		seq_printf(
-			s,
-			"  <%s> die %u: iova=0x%llx pages=%llu %s offset=0x%llx",
-			dmap->dmabufs[0]->exp_name, map->die_index,
-			map->device_address, dmap->size / PAGE_SIZE,
-			edgetpu_dma_dir_rw_s(map->dir), dmap->offset);
+		seq_printf(s, "  <%s> die %u: iova=0x%llx pages=%llu %s",
+			   dmap->dmabufs[0]->exp_name, map->die_index, map->device_address,
+			   dmap->size / PAGE_SIZE, edgetpu_dma_dir_rw_s(map->dir));
 
 	edgetpu_device_dram_dmabuf_info_show(dmap->dmabufs[0], s);
 	seq_puts(s, " dma=");
@@ -502,12 +497,11 @@ err_free:
 }
 
 /*
- * Duplicates @sgt in region [@offset, @offset + @size] to @out.
+ * Duplicates @sgt in region [0, @size) to @out.
  * Only duplicates the "page" parts in @sgt, DMA addresses and lengths are not
  * considered.
  */
-static int dup_sgt_in_region(struct sg_table *sgt, u64 offset, u64 size,
-			     struct sg_table *out)
+static int dup_sgt_in_region(struct sg_table *sgt, u64 size, struct sg_table *out)
 {
 	uint n = 0;
 	u64 cur_offset = 0;
@@ -519,9 +513,8 @@ static int dup_sgt_in_region(struct sg_table *sgt, u64 offset, u64 size,
 	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
 		size_t pg_len = sg->length + sg->offset;
 
-		if (offset < cur_offset + pg_len)
-			n++;
-		if (offset + size <= cur_offset + pg_len)
+		n++;
+		if (size <= cur_offset + pg_len)
 			break;
 		cur_offset += pg_len;
 	}
@@ -532,23 +525,16 @@ static int dup_sgt_in_region(struct sg_table *sgt, u64 offset, u64 size,
 	new_sg = out->sgl;
 	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
 		size_t pg_len = sg->length + sg->offset;
+		struct page *page = sg_page(sg);
+		unsigned int len = pg_len;
+		u64 remain_size = size - cur_offset;
 
-		if (offset < cur_offset + pg_len) {
-			struct page *page = sg_page(sg);
-			unsigned int len = pg_len;
-			u64 remain_size = offset + size - cur_offset;
+		if (remain_size < pg_len)
+			len -= pg_len - remain_size;
+		sg_set_page(new_sg, page, len, 0);
+		new_sg = sg_next(new_sg);
 
-			if (cur_offset < offset) {
-				page = nth_page(page, (offset - cur_offset) >>
-							      PAGE_SHIFT);
-				len -= offset - cur_offset;
-			}
-			if (remain_size < pg_len)
-				len -= pg_len - remain_size;
-			sg_set_page(new_sg, page, len, 0);
-			new_sg = sg_next(new_sg);
-		}
-		if (offset + size <= cur_offset + pg_len)
+		if (size <= cur_offset + pg_len)
 			break;
 		cur_offset += pg_len;
 	}
@@ -556,13 +542,12 @@ static int dup_sgt_in_region(struct sg_table *sgt, u64 offset, u64 size,
 }
 
 /*
- * Copy the DMA addresses and lengths in region [@offset, @offset + @size) from
+ * Copy the DMA addresses and lengths in region [0, @size) from
  * @sgt to @out.
  *
  * The DMA addresses will be condensed when possible.
  */
-static void shrink_sgt_dma_in_region(struct sg_table *sgt, u64 offset, u64 size,
-				     struct sg_table *out)
+static void shrink_sgt_dma_in_region(struct sg_table *sgt, u64 size, struct sg_table *out)
 {
 	u64 cur_offset = 0;
 	struct scatterlist *sg, *prv_sg = NULL, *cur_sg;
@@ -571,20 +556,12 @@ static void shrink_sgt_dma_in_region(struct sg_table *sgt, u64 offset, u64 size,
 	out->nents = 0;
 	for (sg = sgt->sgl; sg;
 	     cur_offset += sg_dma_len(sg), sg = sg_next(sg)) {
-		u64 remain_size = offset + size - cur_offset;
+		u64 remain_size = size - cur_offset;
 		dma_addr_t dma;
 		size_t len;
 
-		/* hasn't touched the first covered sg */
-		if (offset >= cur_offset + sg_dma_len(sg))
-			continue;
 		dma = sg_dma_address(sg);
 		len = sg_dma_len(sg);
-		/* offset exceeds current sg */
-		if (offset > cur_offset) {
-			dma += offset - cur_offset;
-			len -= offset - cur_offset;
-		}
 		if (remain_size < sg_dma_len(sg))
 			len -= sg_dma_len(sg) - remain_size;
 		if (prv_sg &&
@@ -603,15 +580,14 @@ static void shrink_sgt_dma_in_region(struct sg_table *sgt, u64 offset, u64 size,
 	}
 }
 
-static int entry_set_shrunk_sgt(struct dmabuf_map_entry *entry, u64 offset,
-				u64 size)
+static int entry_set_shrunk_sgt(struct dmabuf_map_entry *entry, u64 size)
 {
 	int ret;
 
-	ret = dup_sgt_in_region(entry->sgt, offset, size, &entry->shrunk_sgt);
+	ret = dup_sgt_in_region(entry->sgt, size, &entry->shrunk_sgt);
 	if (ret)
 		return ret;
-	shrink_sgt_dma_in_region(entry->sgt, offset, size, &entry->shrunk_sgt);
+	shrink_sgt_dma_in_region(entry->sgt, size, &entry->shrunk_sgt);
 	return 0;
 }
 
@@ -621,10 +597,8 @@ static int entry_set_shrunk_sgt(struct dmabuf_map_entry *entry, u64 offset,
  *
  * Fields of @entry will be set on success.
  */
-static int etdev_attach_dmabuf_to_entry(struct edgetpu_dev *etdev,
-					struct dma_buf *dmabuf,
-					struct dmabuf_map_entry *entry,
-					u64 offset, u64 size,
+static int etdev_attach_dmabuf_to_entry(struct edgetpu_dev *etdev, struct dma_buf *dmabuf,
+					struct dmabuf_map_entry *entry, u64 size,
 					enum dma_data_direction dir)
 {
 	struct dma_buf_attachment *attachment;
@@ -641,7 +615,7 @@ static int etdev_attach_dmabuf_to_entry(struct edgetpu_dev *etdev,
 	}
 	entry->attachment = attachment;
 	entry->sgt = sgt;
-	ret = entry_set_shrunk_sgt(entry, offset, size);
+	ret = entry_set_shrunk_sgt(entry, size);
 	if (ret)
 		goto err_unmap;
 
@@ -663,7 +637,7 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	struct dma_buf *dmabuf;
 	edgetpu_map_flag_t flags = arg->flags;
 	const u64 offset = arg->offset;
-	const u64 size = PAGE_ALIGN(arg->size);
+	u64 size;
 	const enum dma_data_direction dir =
 		edgetpu_host_dma_dir(flags & EDGETPU_MAP_DIR_MASK);
 	struct edgetpu_dev *etdev;
@@ -674,13 +648,14 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	/* 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)
-		return -EINVAL;
+	/* TODO(b/189278468): entirely ignore @offset */
+	if (offset != 0)
+		etdev_warn_ratelimited(group->etdev,
+				       "Non-zero offset for dmabuf mapping is deprecated");
 	dmabuf = dma_buf_get(arg->dmabuf_fd);
 	if (IS_ERR(dmabuf))
 		return PTR_ERR(dmabuf);
-	if (offset + size > dmabuf->size)
+	if (offset >= dmabuf->size)
 		goto err_put;
 
 	mutex_lock(&group->lock);
@@ -698,13 +673,12 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	get_dma_buf(dmabuf);
 	dmap->dmabufs[0] = dmabuf;
 	dmap->offset = offset;
-	dmap->size = size;
+	dmap->size = size = dmabuf->size;
 	if (IS_MIRRORED(flags)) {
 		for (i = 0; i < group->n_clients; i++) {
 			etdev = edgetpu_device_group_nth_etdev(group, i);
-			ret = etdev_attach_dmabuf_to_entry(etdev, dmabuf,
-							   &dmap->entries[i],
-							   offset, size, dir);
+			ret = etdev_attach_dmabuf_to_entry(etdev, dmabuf, &dmap->entries[i], size,
+							   dir);
 			if (ret)
 				goto err_release_map;
 		}
@@ -718,8 +692,7 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 			ret = -EINVAL;
 			goto err_release_map;
 		}
-		ret = etdev_attach_dmabuf_to_entry(
-			etdev, dmabuf, &dmap->entries[0], offset, size, dir);
+		ret = etdev_attach_dmabuf_to_entry(etdev, dmabuf, &dmap->entries[0], size, dir);
 		if (ret)
 			goto err_release_map;
 		ret = etdev_map_dmabuf(etdev, dmap, dir, &tpu_addr);
@@ -727,11 +700,11 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 			goto err_release_map;
 		dmap->map.die_index = arg->die_index;
 	}
-	dmap->map.device_address = tpu_addr;
+	dmap->map.device_address = tpu_addr + offset;
 	ret = edgetpu_mapping_add(&group->dmabuf_mappings, &dmap->map);
 	if (ret)
 		goto err_release_map;
-	arg->device_address = tpu_addr;
+	arg->device_address = dmap->map.device_address;
 	mutex_unlock(&group->lock);
 	dma_buf_put(dmabuf);
 	return 0;
@@ -825,8 +798,7 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 		if (!bmap->dmabufs[i])
 			continue;
 		etdev = edgetpu_device_group_nth_etdev(group, i);
-		ret = etdev_attach_dmabuf_to_entry(etdev, bmap->dmabufs[i],
-						   &bmap->entries[i], 0,
+		ret = etdev_attach_dmabuf_to_entry(etdev, bmap->dmabufs[i], &bmap->entries[i],
 						   bmap->size, dir);
 		if (ret)
 			goto err_release_bmap;
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index db0f296..5b07632 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -1036,11 +1036,18 @@ static const struct file_operations syncfences_ops = {
 	.release = single_release,
 };
 
-static void edgetpu_debugfs_global_setup(void)
+static int edgetpu_debugfs_global_setup(void)
 {
 	edgetpu_debugfs_dir = debugfs_create_dir("edgetpu", NULL);
+	if (IS_ERR(edgetpu_debugfs_dir)) {
+		pr_err(DRIVER_NAME " error creating edgetpu debugfs dir: %ld\n",
+		       PTR_ERR(edgetpu_debugfs_dir));
+		return PTR_ERR(edgetpu_debugfs_dir);
+	}
+
 	debugfs_create_file("syncfences", 0440, edgetpu_debugfs_dir, NULL,
 			    &syncfences_ops);
+	return 0;
 }
 
 int __init edgetpu_fs_init(void)
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index f48fe20..cdd01d9 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -14,9 +14,6 @@
 #include <linux/types.h>
 #include <linux/version.h>
 
-#ifdef CONFIG_ABROLHOS
-#include "abrolhos-platform.h"
-#endif
 #include "edgetpu-internal.h"
 #include "edgetpu-mapping.h"
 #include "edgetpu-mmu.h"
@@ -234,9 +231,6 @@ out:
 /* mmu_info is unused and NULL for IOMMU version, let IOMMU API supply info */
 int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 {
-#ifdef CONFIG_ABROLHOS
-	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
-#endif
 	struct edgetpu_iommu *etiommu;
 	int ret;
 
@@ -270,28 +264,6 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 
 	/* etiommu initialization done */
 	etdev->mmu_cookie = etiommu;
-	/* TODO (b/178571278): remove chipset specific code. */
-#ifdef CONFIG_ABROLHOS
-	if (!abpdev->csr_iova)
-		goto success;
-
-	etdev_dbg(etdev, "Mapping device CSRs: %llX -> %llX (%lu bytes)\n",
-		  abpdev->csr_iova, abpdev->csr_paddr, abpdev->csr_size);
-
-	/* Add an IOMMU translation for the CSR region */
-	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) {
-		etdev_err(etdev, "Unable to map device CSRs into IOMMU\n");
-		edgetpu_unregister_iommu_device_fault_handler(etdev);
-		etdev->mmu_cookie = NULL;
-		goto err_free;
-	}
-
-success:
-#endif
 	return 0;
 
 err_free:
@@ -306,24 +278,12 @@ void edgetpu_mmu_reset(struct edgetpu_dev *etdev)
 
 void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
 {
-#ifdef CONFIG_ABROLHOS
-	struct abrolhos_platform_dev *abpdev = to_abrolhos_dev(etdev);
-#endif
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	int i, ret;
 
 	if (!etiommu)
 		return;
 
-#ifdef CONFIG_ABROLHOS
-	if (abpdev->csr_iova) {
-		edgetpu_mmu_remove_translation(&abpdev->edgetpu_dev,
-					       abpdev->csr_iova,
-					       abpdev->csr_size,
-					       EDGETPU_CONTEXT_KCI);
-	}
-	abpdev->csr_iova = 0;
-#endif
 	ret = edgetpu_unregister_iommu_device_fault_handler(etdev);
 	if (ret)
 		etdev_warn(etdev,
@@ -364,7 +324,7 @@ static int get_iommu_map_params(struct edgetpu_dev *etdev,
 {
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	size_t size = 0;
-	int prot = __dma_dir_to_iommu_prot(map->dir);
+	int prot = __dma_dir_to_iommu_prot(map->dir, etdev->dev);
 	struct iommu_domain *domain;
 	int i;
 	struct scatterlist *sg;
@@ -461,7 +421,7 @@ int edgetpu_mmu_map_iova_sgt(struct edgetpu_dev *etdev, tpu_addr_t iova,
 			     struct sg_table *sgt, enum dma_data_direction dir,
 			     enum edgetpu_context_id context_id)
 {
-	const int prot = __dma_dir_to_iommu_prot(edgetpu_host_dma_dir(dir));
+	const int prot = __dma_dir_to_iommu_prot(edgetpu_host_dma_dir(dir), etdev->dev);
 	const tpu_addr_t orig_iova = iova;
 	struct scatterlist *sg;
 	int i;
@@ -550,7 +510,7 @@ tpu_addr_t edgetpu_mmu_tpu_map(struct edgetpu_dev *etdev, dma_addr_t down_addr,
 	struct iommu_domain *default_domain =
 		iommu_get_domain_for_dev(etdev->dev);
 	phys_addr_t paddr;
-	int prot = __dma_dir_to_iommu_prot(dir);
+	int prot = __dma_dir_to_iommu_prot(dir, etdev->dev);
 
 	domain = get_domain_by_context_id(etdev, context_id);
 	/*
@@ -602,7 +562,7 @@ tpu_addr_t edgetpu_mmu_tpu_map_sgt(struct edgetpu_dev *etdev,
 	phys_addr_t paddr;
 	dma_addr_t iova, cur_iova;
 	size_t size;
-	int prot = __dma_dir_to_iommu_prot(dir);
+	int prot = __dma_dir_to_iommu_prot(dir, etdev->dev);
 	struct scatterlist *sg;
 	int ret;
 	int i;
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 210d273..1258cf0 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -366,7 +366,13 @@ irqreturn_t edgetpu_chip_irq_handler(int irq, void *arg);
  *
  * Returns 0 on success, otherwise -errno.
  */
-int edgetpu_setup_mmu(struct edgetpu_dev *etdev);
+int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev);
+
+/*
+ * Reverts edgetpu_chip_setup_mmu().
+ * This is called during device removal.
+ */
+void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev);
 
 /* Read TPU timestamp */
 u64 edgetpu_chip_tpu_timestamp(struct edgetpu_dev *etdev);
@@ -387,6 +393,9 @@ struct edgetpu_client *edgetpu_client_add(struct edgetpu_dev *etdev);
 /* Remove TPU client */
 void edgetpu_client_remove(struct edgetpu_client *client);
 
+/* Handle chip-specific client removal */
+void edgetpu_chip_client_remove(struct edgetpu_client *client);
+
 /* mmap() device/queue memory */
 int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma);
 
diff --git a/drivers/edgetpu/edgetpu-mapping.h b/drivers/edgetpu/edgetpu-mapping.h
index 6681427..e3a0dc9 100644
--- a/drivers/edgetpu/edgetpu-mapping.h
+++ b/drivers/edgetpu/edgetpu-mapping.h
@@ -7,6 +7,7 @@
 #ifndef __EDGETPU_MAPPING_H__
 #define __EDGETPU_MAPPING_H__
 
+#include <linux/device.h>
 #include <linux/dma-direction.h>
 #include <linux/iommu.h>
 #include <linux/mutex.h>
@@ -14,6 +15,10 @@
 #include <linux/scatterlist.h>
 #include <linux/seq_file.h>
 #include <linux/types.h>
+#include <linux/version.h>
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 10, 0)
+#include <linux/dma-map-ops.h>
+#endif
 
 #include "edgetpu-internal.h"
 
@@ -139,13 +144,14 @@ void edgetpu_mapping_clear(struct edgetpu_mapping_root *mappings);
 void edgetpu_mappings_show(struct edgetpu_mapping_root *mappings,
 			   struct seq_file *s);
 
-static inline int __dma_dir_to_iommu_prot(enum dma_data_direction dir)
+static inline int __dma_dir_to_iommu_prot(enum dma_data_direction dir, struct device *dev)
 {
-	int prot = 0;
-
-#ifdef EDGETPU_IS_IO_COHERENT
-	prot = IOMMU_CACHE;
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 10, 0)
+	int prot = dev_is_dma_coherent(dev) ? IOMMU_CACHE : 0;
+#else
+	int prot = 0;	/* hardcode to non-dma-coherent for prior kernels */
 #endif
+
 	switch (dir) {
 	case DMA_BIDIRECTIONAL:
 		return prot | IOMMU_READ | IOMMU_WRITE;
diff --git a/drivers/edgetpu/edgetpu-mmu.h b/drivers/edgetpu/edgetpu-mmu.h
index 8ef2f2c..8c5ae3c 100644
--- a/drivers/edgetpu/edgetpu-mmu.h
+++ b/drivers/edgetpu/edgetpu-mmu.h
@@ -118,6 +118,7 @@ static inline unsigned long map_to_dma_attr(edgetpu_map_flag_t flags, bool map)
 
 int edgetpu_mmu_attach(struct edgetpu_dev *dev, void *mmu_info);
 void edgetpu_mmu_detach(struct edgetpu_dev *dev);
+
 /**
  * Re-attach to previously attached MMU.
  *
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index cf9435c..abc9095 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -144,9 +144,7 @@ static void edgetpu_fw_log(struct edgetpu_telemetry *log)
 			       start);
 		if (entry.length == 0 || entry.length > max_length) {
 			header->head = header->tail;
-#if 0 /* TODO(b/170340226): add me back */
 			etdev_err_ratelimited(etdev, "log queue is corrupted");
-#endif
 			break;
 		}
 		copy_with_wrap(header, buffer, entry.length, queue_size, start);
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index 73ee06c..6c94372 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -653,7 +653,22 @@ static ssize_t fw_thread_stats_show(
 	mutex_unlock(&ustats->usage_stats_lock);
 	return ret;
 }
-static DEVICE_ATTR_RO(fw_thread_stats);
+
+static ssize_t fw_thread_stats_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;
+	int i;
+
+	mutex_lock(&ustats->usage_stats_lock);
+	for (i = 0; i < EDGETPU_FW_THREAD_COUNT; i++)
+		ustats->thread_stack_max[i] = 0;
+	mutex_unlock(&ustats->usage_stats_lock);
+	return count;
+}
+static DEVICE_ATTR_RW(fw_thread_stats);
 
 static struct attribute *usage_stats_dev_attrs[] = {
 	&dev_attr_tpu_usage.attr,
diff --git a/drivers/edgetpu/edgetpu.h b/drivers/edgetpu/edgetpu.h
index eacbe80..93d7afa 100644
--- a/drivers/edgetpu/edgetpu.h
+++ b/drivers/edgetpu/edgetpu.h
@@ -273,13 +273,9 @@ struct edgetpu_sync_ioctl {
 	_IOW(EDGETPU_IOCTL_BASE, 16, struct edgetpu_sync_ioctl)
 
 struct edgetpu_map_dmabuf_ioctl {
-	/*
-	 * Offset within the dma-buf to be mapped in bytes.
-	 *
-	 * Must be page-aligned.
-	 */
+	/* Deprecated; pass 0 to keep compatibility. */
 	__u64 offset;
-	/* Size to be mapped in bytes. */
+	/* Ignored; the entire dma-buf is mapped. */
 	__u64 size;
 	/*
 	 * Returned TPU VA.
@@ -310,7 +306,6 @@ struct edgetpu_map_dmabuf_ioctl {
  * On success, @device_address is set and the syscall returns zero.
  *
  * EINVAL: If @offset is not page-aligned.
- * EINVAL: If @size is zero.
  * EINVAL: (for EDGETPU_MAP_NONMIRRORED case) If @die_index exceeds the number
  *         of clients in the group.
  * EINVAL: If the target device group is disbanded.
@@ -321,9 +316,8 @@ struct edgetpu_map_dmabuf_ioctl {
  * Un-map address previously mapped by EDGETPU_MAP_DMABUF.
  *
  * Only fields @die_index and @device_address in the third argument will be
- * used, other fields such as @size and @offset will be fetched from the
- * kernel's internal records. If the buffer was requested as
- * EDGETPU_MAP_MIRRORED, @die_index is ignored as well.
+ * used, other fields will be fetched from the kernel's internal records. If the
+ * buffer was requested as EDGETPU_MAP_MIRRORED, @die_index is ignored as well.
  *
  * EINVAL: If @device_address is not found.
  * EINVAL: If the target device group is disbanded.