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

gcip: check block power state before fetching resp
Bug: 283659596
gcip: Set .is_block_off to gcip_mailbox
Bug: 283659596 (repeat)
gcip: Add gcip_pm_put_async
gcip: iommu: legacy mode do not double map for default domain
gcip: keep track of whether a gdomain is for the default domain
gcip: Add .is_block_off to gcip_{kci,mailbox}
Bug: 283659596 (repeat)
gcip: Add gcip_pm_put_async
gcip: include: keep track of whether a gdomain is for the default domain
edgetpu: Set .is_block_off for gcip_kci
Bug: 283659596 (repeat)
edgetpu: Move pmu_status to edgetpu_dev
Bug: 283659596 (repeat)
edgetpu: hold mmap_lock around pin_user_pages
mmap_lock must be held around call to pin_user_pages to protect against
concurrent address space modifications.
Bug: 283787360
(cherry picked from commit 38761382e80a3446ded8e60e121d5357bdceb8be)
(cherry picked from commit 97d6d45f552d5ee711c15b5085b2af5e726a17ec)
edgetpu: remove code for no KCI data structure
edgetpu: debug prints for iommu mapping calls
edgetpu: Use gcip_iommu_domain_pool_enable_best_fit_algo
Bug: 275112261
edgetpu: Remove the edgetpu_mmu_tpu_map interface
edgetpu: Use per-domain IOVA for coherent buffers
Bug: 275112261 (repeat)
edgetpu: Add edgetpu_mmu_(un)map_sgt()
Bug: 275112261 (repeat)
edgetpu: Use per-domain IOVA for dma-bufs
Bug: 275112261 (repeat)
edgetpu: Adopt gcip-iommu with split IOVA-space
Bug: 275112261 (repeat)
edgetpu: Store a pointer to the default domain
Bug: 275112261 (repeat)
edgetpu: Remove dead mmu_alloc/reserve/free code
gcip: Call iova_domain_init_rcaches on kernel >= 5.18
Bug: 284584586
edgetpu: usage-stats remove unnecessary out of memory log

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

21 files changed, 373 insertions, 596 deletions

diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 18bffa2..d34c951 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -656,29 +656,72 @@ int edgetpu_alloc_coherent(struct edgetpu_dev *etdev, size_t size,
 			   enum edgetpu_context_id context_id)
 {
 	const u32 flags = EDGETPU_MMU_CC_ACCESS | EDGETPU_MMU_HOST | EDGETPU_MMU_COHERENT;
+	int ret;
 
 	mem->vaddr = dma_alloc_coherent(etdev->dev, size, &mem->dma_addr,
 					GFP_KERNEL);
 	if (!mem->vaddr)
 		return -ENOMEM;
+
 	edgetpu_x86_coherent_mem_init(mem);
-	mem->tpu_addr =
-		edgetpu_mmu_tpu_map(etdev, mem->dma_addr, size,
-				    DMA_BIDIRECTIONAL, context_id, flags);
-	if (!mem->tpu_addr) {
-		dma_free_coherent(etdev->dev, size, mem->vaddr, mem->dma_addr);
-		mem->vaddr = NULL;
-		return -EINVAL;
+
+	/* If this context's mappings reside in the default domain, we're done */
+	if (edgetpu_mmu_is_context_using_default_domain(etdev, context_id)) {
+		mem->tpu_addr = mem->dma_addr;
+		mem->size = size;
+		return 0;
+	}
+
+	/*
+	 * dma_get_sgtable may not always be available, and coherent buffers are always physically
+	 * contiguous, so create a 1-entry sgt by hand.
+	 */
+	mem->client_sgt = kzalloc(sizeof(*mem->client_sgt), GFP_KERNEL);
+	if (!mem->client_sgt) {
+		ret = -ENOMEM;
+		goto err_free_coherent;
 	}
+	mem->client_sgt->sgl = kzalloc(sizeof(*mem->client_sgt->sgl), GFP_KERNEL);
+	if (!mem->client_sgt->sgl) {
+		ret = -ENOMEM;
+		goto err_free_sgt;
+	}
+	mem->client_sgt->nents = 1;
+	mem->client_sgt->orig_nents = 1;
+	sg_set_page(mem->client_sgt->sgl, virt_to_page(mem->vaddr), PAGE_ALIGN(size), 0);
+
+	ret = edgetpu_mmu_map_sgt(etdev, mem->client_sgt, context_id, DMA_BIDIRECTIONAL, 0, flags);
+	if (!ret) {
+		etdev_err(etdev, "Failed to map coherent buffer to context %#X\n", context_id);
+		ret = -EIO;
+		goto err_free_sgl;
+	}
+
+	mem->tpu_addr = sg_dma_address(mem->client_sgt->sgl);
 	mem->size = size;
 	return 0;
+
+err_free_sgl:
+	kfree(mem->client_sgt->sgl);
+err_free_sgt:
+	kfree(mem->client_sgt);
+	mem->client_sgt = NULL;
+err_free_coherent:
+	dma_free_coherent(etdev->dev, size, mem->vaddr, mem->dma_addr);
+	mem->vaddr = NULL;
+	return ret;
 }
 
 void edgetpu_free_coherent(struct edgetpu_dev *etdev,
 			   struct edgetpu_coherent_mem *mem,
 			   enum edgetpu_context_id context_id)
 {
-	edgetpu_mmu_tpu_unmap(etdev, mem->tpu_addr, mem->size, context_id);
+	if (!edgetpu_mmu_is_context_using_default_domain(etdev, context_id)) {
+		edgetpu_mmu_unmap_sgt(etdev, mem->client_sgt, context_id, DMA_BIDIRECTIONAL,
+				      /*dma_attrs=*/0, /*mmu_flags=*/0);
+		kfree(mem->client_sgt->sgl);
+		kfree(mem->client_sgt);
+	}
 	edgetpu_x86_coherent_mem_set_wb(mem);
 	dma_free_coherent(etdev->dev, mem->size, mem->vaddr, mem->dma_addr);
 	mem->vaddr = NULL;
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index cb7d9d9..3e94dcd 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -730,8 +730,10 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 		kvfree(pages);
 		return ERR_PTR(-ENOMEM);
 	}
+	mmap_read_lock(current->mm);
 	ret = pin_user_pages(host_addr & PAGE_MASK, num_pages, foll_flags,
 			     pages, vmas);
+	mmap_read_unlock(current->mm);
 	kvfree(vmas);
 	if (ret < 0) {
 		etdev_dbg(etdev, "pin_user_pages failed %u:%pK-%u: %d",
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 4ca8c40..b1b549a 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -33,20 +33,6 @@
 #define to_etfence(gfence) container_of(gfence, struct edgetpu_dma_fence, gfence)
 
 /*
- * Records objects for mapping a dma-buf to an edgetpu_dev.
- */
-struct dmabuf_map_entry {
-	struct dma_buf_attachment *attachment;
-	/* SG table returned by dma_buf_map_attachment() */
-	struct sg_table *sgt;
-	/*
-	 * The SG table that shrunk and condensed from @sgt with region [0, size), where @size is
-	 * the size field in edgetpu_dmabuf_map which owns this entry.
-	 */
-	struct sg_table shrunk_sgt;
-};
-
-/*
  * Records the mapping and other fields needed for mapping a dma-buf to a device
  * group.
  */
@@ -55,7 +41,9 @@ struct edgetpu_dmabuf_map {
 	u64 size; /* size of this mapping in bytes */
 	u32 mmu_flags;
 	struct dma_buf *dmabuf;
-	struct dmabuf_map_entry *map_entry;
+	struct dma_buf_attachment *attachment;
+	/* SG table returned by dma_buf_map_attachment(). Records default domain mapping. */
+	struct sg_table *dma_sgt;
 };
 
 /*
@@ -74,27 +62,17 @@ struct edgetpu_dma_fence {
 static const struct dma_fence_ops edgetpu_dma_fence_ops;
 
 /*
- * Maps @dmap->map_entry.
+ * Maps @dmap->map into its owning device group's context.
  *
  * Caller holds @group->lock.
  */
-static int etdev_map_dmabuf(struct edgetpu_dev *etdev,
-			    struct edgetpu_dmabuf_map *dmap,
-			    tpu_addr_t *tpu_addr_p)
+static int etdev_map_dmabuf(struct edgetpu_dev *etdev, struct edgetpu_dmabuf_map *dmap)
 {
 	struct edgetpu_device_group *group = dmap->map.priv;
 	const enum edgetpu_context_id ctx_id =
 		edgetpu_group_context_id_locked(group);
-	struct dmabuf_map_entry *entry = dmap->map_entry;
-	tpu_addr_t tpu_addr;
-
-	tpu_addr = edgetpu_mmu_tpu_map_sgt(etdev, &entry->shrunk_sgt,
-					   dmap->map.dir,
-					   ctx_id, dmap->mmu_flags);
-	if (!tpu_addr)
-		return -ENOSPC;
-	*tpu_addr_p = tpu_addr;
-	return 0;
+
+	return edgetpu_mmu_map(etdev, &dmap->map, ctx_id, dmap->mmu_flags);
 }
 
 /*
@@ -102,16 +80,13 @@ static int etdev_map_dmabuf(struct edgetpu_dev *etdev,
  *
  * Caller holds @group->lock.
  */
-static void etdev_unmap_dmabuf(struct edgetpu_dev *etdev,
-			       struct edgetpu_dmabuf_map *dmap,
-			       tpu_addr_t tpu_addr)
+static void etdev_unmap_dmabuf(struct edgetpu_dev *etdev, struct edgetpu_dmabuf_map *dmap)
 {
 	struct edgetpu_device_group *group = dmap->map.priv;
 	const enum edgetpu_context_id ctx_id =
 		edgetpu_group_context_id_locked(group);
-	struct dmabuf_map_entry *entry = dmap->map_entry;
 
-	edgetpu_mmu_tpu_unmap_sgt(etdev, tpu_addr, &entry->shrunk_sgt, ctx_id);
+	edgetpu_mmu_unmap(etdev, &dmap->map, ctx_id);
 }
 
 /*
@@ -126,26 +101,22 @@ 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 = map->dir;
-	struct dmabuf_map_entry *entry = dmap->map_entry;
-	tpu_addr_t tpu_addr = map->device_address;
-
-	if (tpu_addr)
-		etdev_unmap_dmabuf(group->etdev, dmap, tpu_addr);
-	sg_free_table(&entry->shrunk_sgt);
-	if (entry->sgt)
-		dma_buf_unmap_attachment(entry->attachment, entry->sgt, dir);
-	if (entry->attachment)
-		dma_buf_detach(dmap->dmabuf, entry->attachment);
+
+	if (dmap->map.device_address)
+		etdev_unmap_dmabuf(group->etdev, dmap);
+	sg_free_table(&dmap->map.sgt);
+	if (dmap->dma_sgt)
+		dma_buf_unmap_attachment(dmap->attachment, dmap->dma_sgt, dir);
+	if (dmap->attachment)
+		dma_buf_detach(dmap->dmabuf, dmap->attachment);
 	dma_buf_put(dmap->dmabuf);
 	edgetpu_device_group_put(group);
-	kfree(dmap->map_entry);
 	kfree(dmap);
 }
 
-static void entry_show_dma_addrs(struct dmabuf_map_entry *entry,
-				 struct seq_file *s)
+static void entry_show_dma_addrs(struct edgetpu_dmabuf_map *dmap, struct seq_file *s)
 {
-	struct sg_table *sgt = &entry->shrunk_sgt;
+	struct sg_table *sgt = &dmap->map.sgt;
 
 	if (sgt->nents == 1) {
 		seq_printf(s, "%pad\n", &sg_dma_address(sgt->sgl));
@@ -164,18 +135,16 @@ static void entry_show_dma_addrs(struct dmabuf_map_entry *entry,
 	}
 }
 
-static void dmabuf_map_callback_show(struct edgetpu_mapping *map,
-				     struct seq_file *s)
+static void dmabuf_map_callback_show(struct edgetpu_mapping *map, struct seq_file *s)
 {
 	struct edgetpu_dmabuf_map *dmap =
 		container_of(map, struct edgetpu_dmabuf_map, map);
 
-	seq_printf(s, "  <%s> iova=%#llx pages=%llu %s",
-		   dmap->dmabuf->exp_name, map->device_address,
-		   DIV_ROUND_UP(dmap->size, PAGE_SIZE),
-		   edgetpu_dma_dir_rw_s(map->dir));
+	seq_printf(s, "  <%s> iova=%#llx (default=%#llx) pages=%llu %s", dmap->dmabuf->exp_name,
+		   map->device_address, sg_dma_address(dmap->dma_sgt->sgl),
+		   DIV_ROUND_UP(dmap->size, PAGE_SIZE), edgetpu_dma_dir_rw_s(map->dir));
 	seq_puts(s, " dma=");
-	entry_show_dma_addrs(dmap->map_entry, s);
+	entry_show_dma_addrs(dmap, s);
 }
 
 /*
@@ -193,9 +162,6 @@ alloc_dmabuf_map(struct edgetpu_device_group *group, edgetpu_map_flag_t flags)
 
 	if (!dmap)
 		return NULL;
-	dmap->map_entry = kzalloc(sizeof(*dmap->map_entry), GFP_KERNEL);
-	if (!dmap->map_entry)
-		goto err_free;
 	dmap->mmu_flags = map_to_mmu_flags(flags) | EDGETPU_MMU_DMABUF;
 	map = &dmap->map;
 	map->flags = flags;
@@ -204,11 +170,6 @@ alloc_dmabuf_map(struct edgetpu_device_group *group, edgetpu_map_flag_t flags)
 	map->show = dmabuf_map_callback_show;
 	map->priv = edgetpu_device_group_get(group);
 	return dmap;
-
-err_free:
-	kfree(dmap->map_entry);
-	kfree(dmap);
-	return NULL;
 }
 
 /*
@@ -257,64 +218,14 @@ static int dup_sgt_in_region(struct sg_table *sgt, u64 size, struct sg_table *ou
 }
 
 /*
- * 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 size, struct sg_table *out)
-{
-	u64 cur_offset = 0;
-	struct scatterlist *sg, *prv_sg = NULL, *cur_sg;
-
-	cur_sg = out->sgl;
-	out->nents = 0;
-	for (sg = sgt->sgl; sg;
-	     cur_offset += sg_dma_len(sg), sg = sg_next(sg)) {
-		u64 remain_size = size - cur_offset;
-		dma_addr_t dma;
-		size_t len;
-
-		dma = sg_dma_address(sg);
-		len = sg_dma_len(sg);
-		if (remain_size < sg_dma_len(sg))
-			len -= sg_dma_len(sg) - remain_size;
-		if (prv_sg &&
-		    sg_dma_address(prv_sg) + sg_dma_len(prv_sg) == dma) {
-			/* merge to previous sg */
-			sg_dma_len(prv_sg) += len;
-		} else {
-			sg_dma_address(cur_sg) = dma;
-			sg_dma_len(cur_sg) = len;
-			prv_sg = cur_sg;
-			cur_sg = sg_next(cur_sg);
-			out->nents++;
-		}
-		if (remain_size <= sg_dma_len(sg))
-			break;
-	}
-}
-
-static int entry_set_shrunk_sgt(struct dmabuf_map_entry *entry, u64 size)
-{
-	int ret;
-
-	ret = dup_sgt_in_region(entry->sgt, size, &entry->shrunk_sgt);
-	if (ret)
-		return ret;
-	shrink_sgt_dma_in_region(entry->sgt, size, &entry->shrunk_sgt);
-	return 0;
-}
-
-/*
  * Performs dma_buf_attach + dma_buf_map_attachment of @dmabuf to @etdev, and
- * sets @entry per the attaching result.
+ * sets @dmap per the attaching result.
  *
- * Fields of @entry will be set on success.
+ * Fields of @dmap 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 size,
-					enum dma_data_direction dir)
+static int etdev_attach_dmabuf_to_mapping(struct edgetpu_dev *etdev, struct dma_buf *dmabuf,
+					  struct edgetpu_dmabuf_map *dmap, u64 size,
+					  enum dma_data_direction dir)
 {
 	struct dma_buf_attachment *attachment;
 	struct sg_table *sgt;
@@ -328,9 +239,9 @@ static int etdev_attach_dmabuf_to_entry(struct edgetpu_dev *etdev, struct dma_bu
 		ret = PTR_ERR(sgt);
 		goto err_detach;
 	}
-	entry->attachment = attachment;
-	entry->sgt = sgt;
-	ret = entry_set_shrunk_sgt(entry, size);
+	dmap->attachment = attachment;
+	dmap->dma_sgt = sgt;
+	ret = dup_sgt_in_region(dmap->dma_sgt, size, &dmap->map.sgt);
 	if (ret)
 		goto err_unmap;
 
@@ -340,8 +251,8 @@ err_unmap:
 	dma_buf_unmap_attachment(attachment, sgt, dir);
 err_detach:
 	dma_buf_detach(dmabuf, attachment);
-	entry->sgt = NULL;
-	entry->attachment = NULL;
+	dmap->dma_sgt = NULL;
+	dmap->attachment = NULL;
 	return ret;
 }
 
@@ -354,7 +265,6 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	u64 size;
 	const enum dma_data_direction dir = map_flag_to_host_dma_dir(flags);
 	struct edgetpu_dmabuf_map *dmap;
-	tpu_addr_t tpu_addr;
 
 	if (!valid_dma_direction(dir)) {
 		etdev_dbg(group->etdev, "%s: invalid direction %d\n", __func__, dir);
@@ -385,35 +295,32 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 	get_dma_buf(dmabuf);
 	dmap->dmabuf = dmabuf;
 	dmap->map.map_size = dmap->size = size = dmabuf->size;
-	ret = etdev_attach_dmabuf_to_entry(group->etdev, dmabuf,
-					   dmap->map_entry, size, dir);
+	ret = etdev_attach_dmabuf_to_mapping(group->etdev, dmabuf, dmap, size, dir);
 	if (ret) {
 		etdev_dbg(group->etdev,
 			  "%s: etdev_attach_dmabuf_to_entry returns %d\n",
 			  __func__, ret);
 		goto err_release_map;
 	}
-	ret = etdev_map_dmabuf(group->etdev, dmap, &tpu_addr);
+	ret = etdev_map_dmabuf(group->etdev, dmap);
 	if (ret) {
 		etdev_dbg(group->etdev,
 			  "%s: etdev_map_dmabuf returns %d\n",
 			  __func__, ret);
 		goto err_release_map;
 	}
-	dmap->map.device_address = tpu_addr;
 	ret = edgetpu_mapping_add(&group->dmabuf_mappings, &dmap->map);
 	if (ret) {
 		etdev_dbg(group->etdev, "%s: edgetpu_mapping_add returns %d\n",
 			  __func__, 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;
 
 err_release_map:
-	/* also releases map_entry if set */
 	dmabuf_map_callback_release(&dmap->map);
 err_unlock_group:
 	mutex_unlock(&group->lock);
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index 1aee8df..1bc4a4e 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -13,7 +13,7 @@
 #include <linux/slab.h>
 #include <linux/types.h>
 
-#include <gcip/gcip-domain-pool.h>
+#include <gcip/gcip-iommu.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
@@ -30,7 +30,12 @@ struct edgetpu_iommu {
 	 * IOMMU domains currently attached.
 	 * NULL for a slot that doesn't have an attached domain.
 	 */
-	struct iommu_domain *domains[EDGETPU_NCONTEXTS];
+	struct gcip_iommu_domain *gdomains[EDGETPU_NCONTEXTS];
+	/*
+	 * Pointer to the default domain. `domains[0]` will always point to `default_domain`, if
+	 * initialization of this structure is successful.
+	 */
+	struct gcip_iommu_domain *default_gdomain;
 	/*
 	 * Records IDs for all domains currently allocated, to support IOMMU (un)mapping
 	 * when the domain is not attached.
@@ -44,14 +49,7 @@ struct edgetpu_iommu {
 	 * required.
 	 * The implementation will fall back to dynamically allocated domains otherwise.
 	 */
-	struct gcip_domain_pool domain_pool;
-
-};
-
-struct edgetpu_iommu_map_params {
-	int prot;
-	size_t size;
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain_pool domain_pool;
 };
 
 /*
@@ -64,29 +62,26 @@ static uint context_id_to_pasid(enum edgetpu_context_id context_id)
 	return (uint)context_id;
 }
 
-static struct iommu_domain *get_domain_by_token(struct edgetpu_iommu *etiommu,
-						int token)
+static struct gcip_iommu_domain *get_domain_by_token(struct edgetpu_iommu *etiommu, int token)
 {
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 
 	mutex_lock(&etiommu->pool_lock);
-	domain = idr_find(&etiommu->domain_id_pool, token);
+	gdomain = idr_find(&etiommu->domain_id_pool, token);
 	mutex_unlock(&etiommu->pool_lock);
-	return domain;
+	return gdomain;
 }
 
-static struct iommu_domain *
-get_domain_by_context_id(struct edgetpu_dev *etdev,
-			 enum edgetpu_context_id ctx_id)
+static struct gcip_iommu_domain *get_domain_by_context_id(struct edgetpu_dev *etdev,
+							  enum edgetpu_context_id ctx_id)
 {
-	struct iommu_domain *domain = NULL;
-	struct device *dev = etdev->dev;
+	struct gcip_iommu_domain *gdomain = NULL;
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	uint pasid;
 
 	/* always return the default domain when AUX is not supported */
 	if (!etiommu->aux_enabled)
-		return iommu_get_domain_for_dev(dev);
+		return etiommu->default_gdomain;
 	if (ctx_id == EDGETPU_CONTEXT_INVALID)
 		return NULL;
 	if (ctx_id & EDGETPU_CONTEXT_DOMAIN_TOKEN)
@@ -94,12 +89,20 @@ get_domain_by_context_id(struct edgetpu_dev *etdev,
 			etiommu, ctx_id ^ EDGETPU_CONTEXT_DOMAIN_TOKEN);
 	pasid = context_id_to_pasid(ctx_id);
 	if (pasid < EDGETPU_NCONTEXTS)
-		domain = etiommu->domains[pasid];
+		gdomain = etiommu->gdomains[pasid];
 
 	/* Fall back to default domain. */
-	if (!domain)
-		domain = iommu_get_domain_for_dev(dev);
-	return domain;
+	if (!gdomain)
+		gdomain = etiommu->default_gdomain;
+	return gdomain;
+}
+
+bool edgetpu_mmu_is_context_using_default_domain(struct edgetpu_dev *etdev,
+						 enum edgetpu_context_id ctx_id)
+{
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+
+	return get_domain_by_context_id(etdev, ctx_id) == etiommu->default_gdomain;
 }
 
 static int edgetpu_iommu_dev_fault_handler(struct iommu_fault *fault,
@@ -143,10 +146,10 @@ static int edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etd
 /* A callback for idr_for_each to release the domains */
 static int edgetpu_idr_free_domain_callback(int id, void *p, void *data)
 {
-	struct iommu_domain *domain = p;
+	struct gcip_iommu_domain *gdomain = p;
 	struct edgetpu_iommu *etiommu = data;
 
-	gcip_domain_pool_free(&etiommu->domain_pool, domain);
+	gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, gdomain);
 	return 0;
 }
 
@@ -154,8 +157,7 @@ static int edgetpu_iommu_fault_handler(struct iommu_domain *domain,
 				       struct device *dev, unsigned long iova,
 				       int flags, void *token)
 {
-	struct edgetpu_iommu_domain *etdomain =
-		(struct edgetpu_iommu_domain *)token;
+	struct edgetpu_iommu_domain *etdomain = (struct edgetpu_iommu_domain *)token;
 
 	dev_dbg(dev, "IOMMU fault on address %08lX. PASID = %u flags = %08X",
 		iova, etdomain->pasid, flags);
@@ -164,13 +166,12 @@ static int edgetpu_iommu_fault_handler(struct iommu_domain *domain,
 }
 
 static void edgetpu_init_etdomain(struct edgetpu_iommu_domain *etdomain,
-				  struct iommu_domain *domain,
-				  int token)
+				  struct gcip_iommu_domain *gdomain, int token)
 {
-	etdomain->iommu_domain = domain;
+	etdomain->gdomain = gdomain;
 	etdomain->pasid = IOMMU_PASID_INVALID;
 	etdomain->token = token;
-	iommu_set_fault_handler(domain, edgetpu_iommu_fault_handler, etdomain);
+	iommu_set_fault_handler(gdomain->domain, edgetpu_iommu_fault_handler, etdomain);
 }
 
 /*
@@ -180,13 +181,13 @@ static void edgetpu_init_etdomain(struct edgetpu_iommu_domain *etdomain,
 static int check_default_domain(struct edgetpu_dev *etdev,
 				struct edgetpu_iommu *etiommu)
 {
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 	int ret;
 	uint pasid;
 
-	domain = iommu_get_domain_for_dev(etdev->dev);
+	gdomain = gcip_iommu_get_domain_for_dev(etdev->dev);
 	/* if default domain exists then we are done */
-	if (domain) {
+	if (gdomain) {
 		etiommu->context_0_default = true;
 		goto out;
 	}
@@ -195,28 +196,29 @@ static int check_default_domain(struct edgetpu_dev *etdev,
 	if (!etiommu->aux_enabled)
 		return -EINVAL;
 
-	domain = gcip_domain_pool_alloc(&etiommu->domain_pool);
-	if (!domain) {
+	gdomain = gcip_iommu_domain_pool_alloc_domain(&etiommu->domain_pool);
+	if (!gdomain) {
 		etdev_warn(etdev, "iommu domain alloc failed");
 		return -EINVAL;
 	}
-	ret = iommu_aux_attach_device(domain, etdev->dev);
+	ret = iommu_aux_attach_device(gdomain->domain, etdev->dev);
 	if (ret) {
 		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
-		gcip_domain_pool_free(&etiommu->domain_pool, domain);
+		gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, gdomain);
 		return ret;
 	}
-	pasid = iommu_aux_get_pasid(domain, etdev->dev);
+	pasid = iommu_aux_get_pasid(gdomain->domain, etdev->dev);
 	/* the default domain must have pasid = 0 */
 	if (pasid != 0) {
 		etdev_warn(etdev, "Invalid PASID %d returned from iommu\n",
 			   pasid);
-		iommu_aux_detach_device(domain, etdev->dev);
-		gcip_domain_pool_free(&etiommu->domain_pool, domain);
+		iommu_aux_detach_device(gdomain->domain, etdev->dev);
+		gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, gdomain);
 		return -EINVAL;
 	}
 out:
-	etiommu->domains[0] = domain;
+	etiommu->default_gdomain = gdomain;
+	etiommu->gdomains[0] = gdomain;
 	return 0;
 }
 
@@ -229,8 +231,19 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	etiommu = kzalloc(sizeof(*etiommu), GFP_KERNEL);
 	if (!etiommu)
 		return -ENOMEM;
-	ret = gcip_domain_pool_init(etdev->dev, &etiommu->domain_pool,
-				       EDGETPU_NUM_PREALLOCATED_DOMAINS);
+	/*
+	 * Specify `base_addr` and `iova_space_size` as 0 so the gcip_iommu_domain_pool will obtain
+	 * the values from the device tree.
+	 */
+	ret = gcip_iommu_domain_pool_init(&etiommu->domain_pool, etdev->dev, 0, 0, SZ_4K,
+					  EDGETPU_NUM_PREALLOCATED_DOMAINS,
+					  GCIP_IOMMU_DOMAIN_TYPE_IOVAD);
+	if (ret) {
+		etdev_err(etdev, "Unable create domain pool (%d)\n", ret);
+		goto err_free_etiommu;
+	}
+	gcip_iommu_domain_pool_enable_best_fit_algo(&etiommu->domain_pool);
+
 	idr_init(&etiommu->domain_id_pool);
 	mutex_init(&etiommu->pool_lock);
 	etiommu->iommu_group = iommu_group_get(etdev->dev);
@@ -246,7 +259,7 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 		etiommu->aux_enabled = true;
 	ret = check_default_domain(etdev, etiommu);
 	if (ret)
-		goto err_free;
+		goto err_destroy_pool;
 
 	ret = edgetpu_register_iommu_device_fault_handler(etdev);
 	if (ret)
@@ -257,7 +270,9 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	etdev->mmu_cookie = etiommu;
 	return 0;
 
-err_free:
+err_destroy_pool:
+	gcip_iommu_domain_pool_destroy(&etiommu->domain_pool);
+err_free_etiommu:
 	kfree(etiommu);
 	return ret;
 }
@@ -277,117 +292,104 @@ void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
 			   ret);
 	for (i = etiommu->context_0_default ? 1 : 0; i < EDGETPU_NCONTEXTS;
 	     i++) {
-		if (etiommu->domains[i])
-			iommu_aux_detach_device(etiommu->domains[i],
-						etdev->dev);
+		if (etiommu->gdomains[i])
+			iommu_aux_detach_device(etiommu->gdomains[i]->domain, etdev->dev);
 	}
 
 	if (etiommu->iommu_group)
 		iommu_group_put(etiommu->iommu_group);
 
 	/* free the domain if the context 0 domain is not default */
-	if (!etiommu->context_0_default && etiommu->domains[0])
-		gcip_domain_pool_free(&etiommu->domain_pool, etiommu->domains[0]);
+	if (!etiommu->context_0_default && etiommu->gdomains[0])
+		gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, etiommu->gdomains[0]);
 
 	idr_for_each(&etiommu->domain_id_pool, edgetpu_idr_free_domain_callback,
 		     etiommu);
 	idr_destroy(&etiommu->domain_id_pool);
-	gcip_domain_pool_destroy(&etiommu->domain_pool);
+	gcip_iommu_domain_pool_destroy(&etiommu->domain_pool);
 	kfree(etiommu);
 	etdev->mmu_cookie = NULL;
 }
 
-static int get_iommu_map_params(struct edgetpu_dev *etdev,
-				struct edgetpu_mapping *map,
-				enum edgetpu_context_id context_id,
-				struct edgetpu_iommu_map_params *params, u32 mmu_flags)
+int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
+		    enum edgetpu_context_id context_id, u32 mmu_flags)
 {
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	size_t size = 0;
-	int prot = mmu_flag_to_iommu_prot(mmu_flags, etdev->dev, map->dir);
-	struct iommu_domain *domain;
-	int i;
-	struct scatterlist *sg;
+	int ret;
 
-	if (!etiommu)
+	ret = edgetpu_mmu_map_sgt(etdev, &map->sgt, context_id, map->dir, map->dma_attrs,
+				  mmu_flags);
+	if (!ret)
 		return -EINVAL;
 
-	domain = get_domain_by_context_id(etdev, context_id);
-	if (!domain) {
-		etdev_err(etdev, "Unable to find an iommu domain\n");
-		return -ENODEV;
-	}
-
-	for_each_sg(map->sgt.sgl, sg, map->sgt.orig_nents, i)
-		size += sg->length;
-
-	prot |= IOMMU_PBHA_PROT(EDGEPTU_MAP_PBHA_VALUE(map->flags));
-	params->prot = prot;
-	params->size = size;
-	params->domain = domain;
+	map->sgt.nents = ret;
+	map->device_address = sg_dma_address(map->sgt.sgl);
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#llx dma=%#llx size=%#x flags=%#x\n",
+		  __func__, context_id, map->device_address, (u64)sg_dma_address(map->sgt.sgl),
+		  sg_dma_len(map->sgt.sgl), mmu_flags);
 	return 0;
 }
 
-int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
-		    enum edgetpu_context_id context_id, u32 mmu_flags)
+void edgetpu_mmu_unmap(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
+		       enum edgetpu_context_id context_id)
 {
+	edgetpu_mmu_unmap_sgt(etdev, &map->sgt, context_id, map->dir, map->dma_attrs, 0);
+}
+
+int edgetpu_mmu_map_sgt(struct edgetpu_dev *etdev, struct sg_table *sgt,
+			enum edgetpu_context_id context_id, enum dma_data_direction dir,
+			unsigned long dma_attrs, u32 mmu_flags)
+
+{
+	struct gcip_iommu_domain *gdomain;
+	u64 gcip_map_flags =
+		GCIP_MAP_FLAGS_DMA_DIRECTION_TO_FLAGS(dir) |
+		GCIP_MAP_FLAGS_DMA_COHERENT_TO_FLAGS((mmu_flags & EDGETPU_MMU_COHERENT) != 0) |
+		GCIP_MAP_FLAGS_DMA_ATTR_TO_FLAGS(dma_attrs);
 	int ret;
-	unsigned long iova;
-	struct edgetpu_iommu_map_params params;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
 
-	ret = get_iommu_map_params(etdev, map, context_id, &params, mmu_flags);
-	if (ret)
-		return ret;
+	gdomain = get_domain_by_context_id(etdev, context_id);
+	if (!gdomain) {
+		etdev_err(etdev, "Unable to find an iommu_domain for context_id %#x\n", context_id);
+		return 0;
+	}
 
-	ret = dma_map_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.nents, map->dir, map->dma_attrs);
+	ret = gcip_iommu_domain_map_sg(gdomain, sgt->sgl, sgt->nents, gcip_map_flags);
 	if (!ret)
-		return -EINVAL;
-	map->sgt.nents = ret;
-	iova = sg_dma_address(map->sgt.sgl);
+		return 0;
 
-	/*
-	 * All mappings get added to the default domain by the call to
-	 * dma_map_sg above.
-	 * Per-context mappings are mirrored to their specific domains here
-	 */
-	if (params.domain != default_domain) {
-		ssize_t mapped = (ssize_t)iommu_map_sg(params.domain, iova, map->sgt.sgl,
-						       map->sgt.orig_nents, params.prot);
-
-		/* iommu_map_sg returns 0 on failure before 5.15, returns -errno afterwards */
-		if (mapped <= 0) {
-			/* Undo the mapping in the default domain */
-			dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.orig_nents, map->dir,
-					   DMA_ATTR_SKIP_CPU_SYNC);
-			return mapped == 0 ? -ENOMEM : (int)mapped;
-		}
-	}
+	/* TODO(b/281157263): Remove once gcip-iommu checks DMA_ATTR_SKIP_CPU_SYNC */
+	if (!(dma_attrs & DMA_ATTR_SKIP_CPU_SYNC))
+		dma_sync_sg_for_device(etdev->dev, sgt->sgl, sgt->orig_nents, dir);
 
-	map->device_address = iova;
-	return 0;
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#llx size=%#x flags=%#llx\n",
+		  __func__, context_id, (u64)sg_dma_address(sgt->sgl), sg_dma_len(sgt->sgl),
+		  gcip_map_flags);
+	return ret;
 }
 
-void edgetpu_mmu_unmap(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
-		       enum edgetpu_context_id context_id)
+void edgetpu_mmu_unmap_sgt(struct edgetpu_dev *etdev, struct sg_table *sgt,
+			   enum edgetpu_context_id context_id, enum dma_data_direction dir,
+			   unsigned long dma_attrs, u32 mmu_flags)
 {
-	int ret;
-	struct edgetpu_iommu_map_params params;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
-
-	ret = get_iommu_map_params(etdev, map, context_id, &params, 0);
-	if (!ret && params.domain != default_domain) {
-		/*
-		 * If this is a per-context mapping, it was mirrored in the
-		 * per-context domain. Undo that mapping first.
-		 */
-		iommu_unmap(params.domain, map->device_address, params.size);
+	struct gcip_iommu_domain *gdomain;
+	u64 gcip_map_flags =
+		GCIP_MAP_FLAGS_DMA_DIRECTION_TO_FLAGS(dir) |
+		GCIP_MAP_FLAGS_DMA_COHERENT_TO_FLAGS((mmu_flags & EDGETPU_MMU_COHERENT) != 0) |
+		GCIP_MAP_FLAGS_DMA_ATTR_TO_FLAGS(dma_attrs);
+
+	gdomain = get_domain_by_context_id(etdev, context_id);
+	if (!gdomain) {
+		etdev_err(etdev, "Unable to find an iommu_domain\n");
+		return;
 	}
 
-	/* Undo the mapping in the default domain */
-	dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.orig_nents, map->dir, map->dma_attrs);
+	/* TODO(b/281157263): Remove once gcip-iommu checks DMA_ATTR_SKIP_CPU_SYNC */
+	if (!(dma_attrs & DMA_ATTR_SKIP_CPU_SYNC))
+		dma_sync_sg_for_cpu(etdev->dev, sgt->sgl, sgt->orig_nents, dir);
+
+	gcip_iommu_domain_unmap_sg(gdomain, sgt->sgl, sgt->orig_nents, gcip_map_flags);
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#llx size=%#x\n",
+		  __func__, context_id, (u64)sg_dma_address(sgt->sgl), sg_dma_len(sgt->sgl));
 }
 
 int edgetpu_mmu_map_iova_sgt(struct edgetpu_dev *etdev, tpu_addr_t iova,
@@ -408,6 +410,8 @@ int edgetpu_mmu_map_iova_sgt(struct edgetpu_dev *etdev, tpu_addr_t iova,
 			goto error;
 		iova += sg->length;
 	}
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#llx dma=%#llx size=%#llx dir=%d\n",
+		  __func__, context_id, iova, (u64)sg_dma_address(sgt->sgl), orig_iova - iova, dir);
 	return 0;
 
 error:
@@ -428,167 +432,34 @@ void edgetpu_mmu_unmap_iova_sgt_attrs(struct edgetpu_dev *etdev,
 
 	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i)
 		size += sg->length;
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#llx size=%#zx\n", __func__, context_id, iova, size);
 	edgetpu_mmu_remove_translation(etdev, iova, size, context_id);
 }
 
-tpu_addr_t edgetpu_mmu_alloc(struct edgetpu_dev *etdev, size_t size,
-			     u32 mmu_flags)
-{
-	return 0;
-}
-
-void edgetpu_mmu_reserve(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-			 size_t size)
-{
-}
-
-void edgetpu_mmu_free(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-		      size_t size)
-{
-}
-
 int edgetpu_mmu_add_translation(struct edgetpu_dev *etdev, unsigned long iova,
 				phys_addr_t paddr, size_t size, int prot,
 				enum edgetpu_context_id context_id)
 {
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 
-	domain = get_domain_by_context_id(etdev, context_id);
-	if (!domain)
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#lx paddr=%pap size=%#zx prot=%#x\n",
+		  __func__, context_id, iova, &paddr, size, prot);
+	gdomain = get_domain_by_context_id(etdev, context_id);
+	if (!gdomain)
 		return -ENODEV;
-	return iommu_map(domain, iova, paddr, size, prot);
+	return iommu_map(gdomain->domain, iova, paddr, size, prot);
 }
 
 void edgetpu_mmu_remove_translation(struct edgetpu_dev *etdev,
 				    unsigned long iova, size_t size,
 				    enum edgetpu_context_id context_id)
 {
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 
-	domain = get_domain_by_context_id(etdev, context_id);
-	if (domain)
-		iommu_unmap(domain, iova, size);
-}
-
-/*
- * This function assumes [@down_addr, @down_addr + size) is mapped to
- * [phys_addr, phys_addr + size). This is true if @down_addr was mapped by
- * dma_alloc_* series, and may not be true when mapped by dma_map_sg*.
- */
-tpu_addr_t edgetpu_mmu_tpu_map(struct edgetpu_dev *etdev, dma_addr_t down_addr,
-			       size_t size, enum dma_data_direction dir,
-			       enum edgetpu_context_id context_id,
-			       u32 mmu_flags)
-{
-	struct iommu_domain *domain;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
-	phys_addr_t paddr;
-	int prot = mmu_flag_to_iommu_prot(mmu_flags, etdev->dev, dir);
-
-	domain = get_domain_by_context_id(etdev, context_id);
-	/*
-	 * Either we don't have per-context domains or this mapping
-	 * belongs to the default context, in which case we don't need
-	 * to do anything
-	 */
-	if (!domain || domain == default_domain)
-		return down_addr;
-	paddr = iommu_iova_to_phys(default_domain, down_addr);
-	if (!paddr)
-		return 0;
-	/* Map the address to the context-specific domain */
-	if (iommu_map(domain, down_addr, paddr, size, prot))
-		return 0;
-
-	/* Return downstream IOMMU DMA address as TPU address. */
-	return down_addr;
-}
-
-void edgetpu_mmu_tpu_unmap(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-			   size_t size, enum edgetpu_context_id context_id)
-{
-	struct iommu_domain *domain;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
-
-	domain = get_domain_by_context_id(etdev, context_id);
-	/*
-	 * Either we don't have per-context domains or this mapping
-	 * belongs to the default context, in which case we don't need
-	 * to do anything
-	 */
-	if (!domain || domain == default_domain)
-		return;
-	/* Unmap the address from the context-specific domain */
-	iommu_unmap(domain, tpu_addr, size);
-}
-
-tpu_addr_t edgetpu_mmu_tpu_map_sgt(struct edgetpu_dev *etdev,
-				   struct sg_table *sgt, enum dma_data_direction dir,
-				   enum edgetpu_context_id context_id,
-				   u32 mmu_flags)
-{
-	struct iommu_domain *domain;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
-	phys_addr_t paddr;
-	dma_addr_t iova, cur_iova;
-	size_t size;
-	int prot = mmu_flag_to_iommu_prot(mmu_flags, etdev->dev, dir);
-	struct scatterlist *sg;
-	int ret;
-	int i;
-
-	/*
-	 * We cannot map the SG to a single TPU VA if the table contains more
-	 * than one DMA address.
-	 */
-	if (sgt->nents != 1)
-		return 0;
-	iova = sg_dma_address(sgt->sgl);
-	domain = get_domain_by_context_id(etdev, context_id);
-	/*
-	 * Either we don't have per-context domains or this mapping
-	 * belongs to the default context, in which case we don't need
-	 * to do anything.
-	 */
-	if (!domain || domain == default_domain)
-		return iova;
-	cur_iova = iova;
-	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
-		/* ignore sg->offset */
-		paddr =  page_to_phys(sg_page(sg));
-		size = sg->length + sg->offset;
-		ret = iommu_map(domain, cur_iova, paddr, size, prot);
-		if (ret)
-			goto rollback;
-		cur_iova += size;
-	}
-
-	return iova;
-rollback:
-	iommu_unmap(domain, iova, cur_iova - iova);
-	etdev_err(etdev, "TPU map sgt failed: %d", ret);
-	return 0;
-}
-
-void edgetpu_mmu_tpu_unmap_sgt(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-			       struct sg_table *sgt,
-			       enum edgetpu_context_id context_id)
-{
-	struct iommu_domain *domain;
-	struct iommu_domain *default_domain =
-		iommu_get_domain_for_dev(etdev->dev);
-
-	domain = get_domain_by_context_id(etdev, context_id);
-	if (!domain || domain == default_domain)
-		return;
-	/*
-	 * We have checked sgt->nents == 1 on map, sg_dma_len(sgt->sgl) should
-	 * equal the total size.
-	 */
-	iommu_unmap(domain, tpu_addr, sg_dma_len(sgt->sgl));
+	etdev_dbg(etdev, "%s: ctx=%x iova=%#lx size=%#zx\n", __func__, context_id, iova, size);
+	gdomain = get_domain_by_context_id(etdev, context_id);
+	if (gdomain)
+		iommu_unmap(gdomain->domain, iova, size);
 }
 
 /* to be returned when domain aux is not supported */
@@ -601,35 +472,35 @@ struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_iommu_domain *etdomain;
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 	int token;
 
 	if (!etiommu->aux_enabled)
 		return &invalid_etdomain;
-	domain = gcip_domain_pool_alloc(&etiommu->domain_pool);
-	if (!domain) {
+	gdomain = gcip_iommu_domain_pool_alloc_domain(&etiommu->domain_pool);
+	if (!gdomain) {
 		etdev_warn(etdev, "iommu domain allocation failed");
 		return NULL;
 	}
 
 	etdomain = kzalloc(sizeof(*etdomain), GFP_KERNEL);
 	if (!etdomain) {
-		gcip_domain_pool_free(&etiommu->domain_pool, domain);
+		gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, gdomain);
 		return NULL;
 	}
 
 	mutex_lock(&etiommu->pool_lock);
-	token = idr_alloc(&etiommu->domain_id_pool, domain, 0,
-			  EDGETPU_DOMAIN_TOKEN_END, GFP_KERNEL);
+	token = idr_alloc(&etiommu->domain_id_pool, gdomain, 0, EDGETPU_DOMAIN_TOKEN_END,
+			  GFP_KERNEL);
 	mutex_unlock(&etiommu->pool_lock);
 	if (token < 0) {
 		etdev_warn(etdev, "alloc iommu domain token failed: %d", token);
 		kfree(etdomain);
-		gcip_domain_pool_free(&etiommu->domain_pool, domain);
+		gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, gdomain);
 		return NULL;
 	}
 
-	edgetpu_init_etdomain(etdomain, domain, token);
+	edgetpu_init_etdomain(etdomain, gdomain, token);
 	return etdomain;
 }
 
@@ -647,7 +518,7 @@ void edgetpu_mmu_free_domain(struct edgetpu_dev *etdev,
 	mutex_lock(&etiommu->pool_lock);
 	idr_remove(&etiommu->domain_id_pool, etdomain->token);
 	mutex_unlock(&etiommu->pool_lock);
-	gcip_domain_pool_free(&etiommu->domain_pool, etdomain->iommu_domain);
+	gcip_iommu_domain_pool_free_domain(&etiommu->domain_pool, etdomain->gdomain);
 	kfree(etdomain);
 }
 
@@ -655,7 +526,7 @@ int edgetpu_mmu_attach_domain(struct edgetpu_dev *etdev,
 			      struct edgetpu_iommu_domain *etdomain)
 {
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	struct iommu_domain *domain;
+	struct gcip_iommu_domain *gdomain;
 	int ret;
 	uint pasid;
 
@@ -664,13 +535,13 @@ int edgetpu_mmu_attach_domain(struct edgetpu_dev *etdev,
 		return 0;
 	if (etdomain->pasid != IOMMU_PASID_INVALID)
 		return -EINVAL;
-	domain = etdomain->iommu_domain;
-	ret = iommu_aux_attach_device(domain, etdev->dev);
+	gdomain = etdomain->gdomain;
+	ret = iommu_aux_attach_device(gdomain->domain, etdev->dev);
 	if (ret) {
 		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
 		return ret;
 	}
-	pasid = iommu_aux_get_pasid(domain, etdev->dev);
+	pasid = iommu_aux_get_pasid(gdomain->domain, etdev->dev);
 	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS) {
 		etdev_warn(etdev, "Invalid PASID %d returned from iommu",
 			   pasid);
@@ -678,15 +549,15 @@ int edgetpu_mmu_attach_domain(struct edgetpu_dev *etdev,
 		goto err_detach;
 	}
 	/* the IOMMU driver returned a duplicate PASID */
-	if (etiommu->domains[pasid]) {
+	if (etiommu->gdomains[pasid]) {
 		ret = -EBUSY;
 		goto err_detach;
 	}
-	etiommu->domains[pasid] = domain;
+	etiommu->gdomains[pasid] = gdomain;
 	etdomain->pasid = pasid;
 	return 0;
 err_detach:
-	iommu_aux_detach_device(domain, etdev->dev);
+	iommu_aux_detach_device(gdomain->domain, etdev->dev);
 	return ret;
 }
 
@@ -700,7 +571,7 @@ void edgetpu_mmu_detach_domain(struct edgetpu_dev *etdev,
 		return;
 	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS)
 		return;
-	etiommu->domains[pasid] = NULL;
+	etiommu->gdomains[pasid] = NULL;
 	etdomain->pasid = IOMMU_PASID_INVALID;
-	iommu_aux_detach_device(etdomain->iommu_domain, etdev->dev);
+	iommu_aux_detach_device(etdomain->gdomain->domain, etdev->dev);
 }
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 17dc48c..e4f17bd 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -73,14 +73,19 @@ typedef u64 tpu_addr_t;
 
 struct edgetpu_coherent_mem {
 	void *vaddr;		/* kernel VA, no allocation if NULL */
-	dma_addr_t dma_addr;	/* DMA handle for downstream IOMMU, if any */
-	tpu_addr_t tpu_addr;	/* DMA handle for TPU internal IOMMU, if any */
+	dma_addr_t dma_addr;	/* IOVA for default domain, returned by dma_alloc_coherent */
+	tpu_addr_t tpu_addr;	/*
+				 * IOVA for the domain of the context the memory was requested for.
+				 * Equal to dma_addr if requested for EDGETPU_CONTEXT_KCI.
+				 */
 	u64 host_addr;		/* address mapped on host for debugging */
 	u64 phys_addr;		/* physical address, if available */
 	size_t size;
 #ifdef CONFIG_X86
 	bool is_set_uc;		/* memory has been marked uncached on X86 */
 #endif
+	/* SGT used to map the coherent memory into the destination context. */
+	struct sg_table *client_sgt;
 };
 
 struct edgetpu_device_group;
@@ -233,6 +238,9 @@ struct edgetpu_dev {
 	/* debug dump handlers */
 	edgetpu_debug_dump_handlers *debug_dump_handlers;
 	struct work_struct debug_dump_work;
+
+	/* PMU status base address for block status, maybe NULL */
+	void __iomem *pmu_status;
 };
 
 struct edgetpu_dev_iface {
diff --git a/drivers/edgetpu/edgetpu-iremap-pool.h b/drivers/edgetpu/edgetpu-iremap-pool.h
index 4024b86..55e6277 100644
--- a/drivers/edgetpu/edgetpu-iremap-pool.h
+++ b/drivers/edgetpu/edgetpu-iremap-pool.h
@@ -27,7 +27,7 @@ void edgetpu_iremap_pool_destroy(struct edgetpu_dev *etdev);
 /*
  * Attempt to allocate memory in the instruction remap pool if the device
  * has one.
- * Fall back to dma_alloc_coherent and edgetpu_mmu_tpu_map otherwise.
+ * Fall back to dma_alloc_coherent and edgetpu_mmu_map_sgt otherwise.
  */
 int edgetpu_iremap_alloc(struct edgetpu_dev *etdev, size_t size,
 			 struct edgetpu_coherent_mem *mem,
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 7a90346..27d0a3f 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -10,6 +10,7 @@
 #include <linux/circ_buf.h>
 #include <linux/device.h>
 #include <linux/errno.h>
+#include <linux/io.h>
 #include <linux/kernel.h>
 #include <linux/seq_file.h>
 #include <linux/slab.h>
@@ -187,6 +188,13 @@ static inline void edgetpu_kci_trigger_doorbell(struct gcip_kci *kci,
 		EDGETPU_MAILBOX_CMD_QUEUE_WRITE(mailbox, doorbell_set, 1);
 }
 
+static inline bool edgetpu_kci_is_block_off(struct gcip_kci *kci)
+{
+	struct edgetpu_mailbox *mailbox = gcip_kci_get_data(kci);
+
+	return mailbox->etdev->pmu_status ? !readl(mailbox->etdev->pmu_status) : false;
+}
+
 static const struct gcip_kci_ops kci_ops = {
 	.get_cmd_queue_head = edgetpu_kci_get_cmd_queue_head,
 	.get_cmd_queue_tail = edgetpu_kci_get_cmd_queue_tail,
@@ -198,6 +206,7 @@ static const struct gcip_kci_ops kci_ops = {
 	.trigger_doorbell = edgetpu_kci_trigger_doorbell,
 	.reverse_kci_handle_response = edgetpu_reverse_kci_handle_response,
 	.update_usage = edgetpu_kci_update_usage_wrapper,
+	.is_block_off = edgetpu_kci_is_block_off,
 };
 
 int edgetpu_kci_init(struct edgetpu_mailbox_manager *mgr, struct edgetpu_kci *etkci)
@@ -537,9 +546,6 @@ void edgetpu_kci_mappings_show(struct edgetpu_dev *etdev, struct seq_file *s)
 	struct edgetpu_coherent_mem *cmd_queue_mem = &etkci->cmd_queue_mem;
 	struct edgetpu_coherent_mem *resp_queue_mem = &etkci->resp_queue_mem;
 
-	if (!etkci || !etkci->mailbox)
-		return;
-
 	seq_printf(s, "kci context mbox %u:\n", EDGETPU_CONTEXT_KCI);
 	seq_printf(s, "  %#llx %lu cmdq - %pad\n", cmd_queue_mem->tpu_addr,
 		   DIV_ROUND_UP(QUEUE_SIZE * gcip_kci_queue_element_size(GCIP_MAILBOX_CMD_QUEUE),
@@ -559,9 +565,6 @@ int edgetpu_kci_shutdown(struct edgetpu_kci *etkci)
 		.code = GCIP_KCI_CODE_SHUTDOWN,
 	};
 
-	if (!etkci || !etkci->kci)
-		return -ENODEV;
-
 	return gcip_kci_send_cmd(etkci->kci, &cmd);
 }
 
@@ -577,9 +580,6 @@ int edgetpu_kci_get_debug_dump(struct edgetpu_kci *etkci, tpu_addr_t tpu_addr, s
 		},
 	};
 
-	if (!etkci || !etkci->kci)
-		return -ENODEV;
-
 	return gcip_kci_send_cmd(etkci->kci, &cmd);
 }
 
@@ -598,9 +598,6 @@ int edgetpu_kci_open_device(struct edgetpu_kci *etkci, u32 mailbox_map, u32 clie
 		},
 	};
 
-	if (!etkci || !etkci->kci)
-		return -ENODEV;
-
 	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(etkci, "open device");
 	if (vcid < 0)
 		return gcip_kci_send_cmd(etkci->kci, &cmd);
@@ -617,9 +614,6 @@ int edgetpu_kci_close_device(struct edgetpu_kci *etkci, u32 mailbox_map)
 		},
 	};
 
-	if (!etkci || !etkci->kci)
-		return -ENODEV;
-
 	RETURN_ERRNO_IF_ETDEV_NOT_GOOD(etkci, "close device");
 
 	return gcip_kci_send_cmd(etkci->kci, &cmd);
@@ -634,9 +628,6 @@ int edgetpu_kci_notify_throttling(struct edgetpu_dev *etdev, u32 level)
 		},
 	};
 
-	if (!etdev->etkci)
-		return -ENODEV;
-
 	return gcip_kci_send_cmd(etdev->etkci->kci, &cmd);
 }
 
@@ -649,9 +640,6 @@ int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block)
 		},
 	};
 
-	if (!etdev->etkci)
-		return -ENODEV;
-
 	return gcip_kci_send_cmd(etdev->etkci->kci, &cmd);
 }
 
@@ -667,9 +655,6 @@ int edgetpu_kci_firmware_tracing_level(void *data, unsigned long level, unsigned
 	struct gcip_kci_response_element resp;
 	int ret;
 
-	if (!etdev->etkci)
-		return -ENODEV;
-
 	ret = gcip_kci_send_cmd_return_resp(etdev->etkci->kci, &cmd, &resp);
 	if (ret == GCIP_KCI_ERROR_OK)
 		*active_level = resp.retval;
@@ -696,9 +681,6 @@ int edgetpu_kci_set_device_properties(struct edgetpu_kci *etkci, struct edgetpu_
 	};
 	int ret = 0;
 
-	if (!etkci || !etkci->kci)
-		return -ENODEV;
-
 	mutex_lock(&dev_prop->lock);
 	if (!dev_prop->initialized)
 		goto out;
@@ -718,8 +700,5 @@ int edgetpu_kci_resp_rkci_ack(struct edgetpu_dev *etdev, struct gcip_kci_respons
 		.code = GCIP_KCI_CODE_RKCI_ACK,
 	};
 
-	if (!etdev->etkci)
-		return -ENODEV;
-
 	return gcip_kci_send_cmd(etdev->etkci->kci, &cmd);
 }
diff --git a/drivers/edgetpu/edgetpu-mmu.h b/drivers/edgetpu/edgetpu-mmu.h
index bff24c6..912aa03 100644
--- a/drivers/edgetpu/edgetpu-mmu.h
+++ b/drivers/edgetpu/edgetpu-mmu.h
@@ -13,6 +13,8 @@
 #include <linux/scatterlist.h>
 #include <linux/version.h>
 
+#include <gcip/gcip-iommu.h>
+
 #include "edgetpu-internal.h"
 #include "edgetpu.h"
 
@@ -39,7 +41,7 @@ struct edgetpu_iommu_domain {
 	 * edgetpu_mmu_detach_domain().
 	 */
 	uint pasid;
-	struct iommu_domain *iommu_domain;
+	struct gcip_iommu_domain *gdomain;
 	/*
 	 * A token set by edgetpu_mmu_alloc_domain(). See the description of
 	 * edgetpu_mmu_add_translation() about @context_id for more details.
@@ -114,6 +116,13 @@ int edgetpu_mmu_map(struct edgetpu_dev *dev, struct edgetpu_mapping *map,
 void edgetpu_mmu_unmap(struct edgetpu_dev *dev, struct edgetpu_mapping *map,
 		       enum edgetpu_context_id context_id);
 
+int edgetpu_mmu_map_sgt(struct edgetpu_dev *etdev, struct sg_table *sgt,
+			enum edgetpu_context_id context_id, enum dma_data_direction dir,
+			unsigned long dma_attrs, u32 mmu_flags);
+void edgetpu_mmu_unmap_sgt(struct edgetpu_dev *etdev, struct sg_table *sgt,
+			   enum edgetpu_context_id context_id, enum dma_data_direction dir,
+			   unsigned long dma_attrs, u32 mmu_flags);
+
 /**
  * Maps TPU IOVA @iova to @sgt.
  * @sgt: the sg table presents the list of pages.
@@ -121,9 +130,6 @@ void edgetpu_mmu_unmap(struct edgetpu_dev *dev, struct edgetpu_mapping *map,
  * Description: Request TPU to map @iova to the pages presented by @sgt.
  *
  * Returns 0 on success, -errno on error.
- *
- * Note: Caller should use edgetpu_mmu_reserve() before calling this method if
- * the target @iova isn't acquired from edgetpu_mmu_alloc(@etdev).
  */
 int edgetpu_mmu_map_iova_sgt(struct edgetpu_dev *etdev, tpu_addr_t iova,
 			     struct sg_table *sgt, enum dma_data_direction dir,
@@ -137,36 +143,6 @@ void edgetpu_mmu_unmap_iova_sgt_attrs(struct edgetpu_dev *etdev,
 	edgetpu_mmu_unmap_iova_sgt_attrs(e, i, s, d, c, 0)
 
 /**
- * Allocates an IOVA in the internal MMU.
- * @size: size needed to be allocated in bytes.
- *
- * Description: Allocates a TPU address to be mapped via
- * edgetpu_mmu_add_translation().
- *
- * If the chip doesn't have an internal MMU then return zero.
- *
- * Returns zero on error.
- */
-tpu_addr_t edgetpu_mmu_alloc(struct edgetpu_dev *etdev, size_t size,
-			     u32 mmu_flags);
-/**
- * Marks the IOVA region [@tpu_addr, @tpu_addr + @size) as reserved.
- *
- * Description: Use this function to mark the region as reserved and prevents
- * it from being allocated by edgetpu_mmu_alloc().
- *
- * Use edgetpu_mmu_free() to release the reserved area.
- */
-void edgetpu_mmu_reserve(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-			 size_t size);
-/*
- * Free the IOVA allocated by edgetpu_mmu_alloc() or reserved by
- * edgetpu_mmu_reserve().
- */
-void edgetpu_mmu_free(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-		      size_t size);
-
-/**
  * Add an IOVA translation to the chip MMU/IOMMU.
  * @iova: I/O virtual address (TPU VA) to map to paddr.
  * @paddr: Physical/next-stage target address to which iova is to be mapped.
@@ -178,10 +154,6 @@ void edgetpu_mmu_free(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
  * paddr can be considered a physical address from the TPU's viewpoint, but
  * may actually be another IOVA for another IOMMU downstream of the chip MMU.
  *
- * Note: for chipsets with edgetpu_mmu_alloc() support, @iova passed to this
- * function must be either allocated from edgetpu_mmu_alloc() or reserved by
- * edgetpu_mmu_reserve().
- *
  * For chipsets with IOMMU AUX domain support, @context_id can be used to
  * specify a detached IOMMU domain by value
  * (EDGETPU_CONTEXT_DOMAIN_TOKEN | @token), where @token is the one returned by
@@ -197,67 +169,6 @@ void edgetpu_mmu_remove_translation(struct edgetpu_dev *etdev,
 				    unsigned long iova, size_t size,
 				    enum edgetpu_context_id context_id);
 
-/**
- * Add a TPU mapping for a local DMA mapping
- * @down_addr: DMA (or physical) addr of memory downstream from TPU
- * @size: size of memory area in bytes
- * @dir: DMA direction of mapping
- * @context_id: context ID for the mapping
- * @mmu_flags: the flag or'ed with EDGETPU_MMU_* macros
- *
- * Description: For chips with internal MMUs, add the required internal MMU
- * mapping for the TPU to access @down_addr, the DMA or physical address of the
- * buffer as returned by the Linux DMA API when the DMA mapping was created.
- * This can be used with, for example, buffers allocated using
- * dma_alloc_coherent(), which are mapped appropriately for any downstream IOMMU
- * and must be mapped to the TPU internal MMU as well.
- *
- * For a chip that doesn't have an internal MMU but has the IOMMU domain AUX
- * feature, perform the necessary mapping to @context_id and return the
- * downstream DMA address.
- *
- * Returns zero on error.
- */
-tpu_addr_t edgetpu_mmu_tpu_map(struct edgetpu_dev *etdev, dma_addr_t down_addr,
-			       size_t size, enum dma_data_direction dir,
-			       enum edgetpu_context_id context_id,
-			       u32 mmu_flags);
-
-/* Unmap a TPU mapping created by edgetpu_mmu_tpu_map */
-void edgetpu_mmu_tpu_unmap(struct edgetpu_dev *etdev,
-			   tpu_addr_t tpu_addr, size_t size,
-			   enum edgetpu_context_id context_id);
-
-/**
- * Add a TPU mapping towards an SG table.
- * @sgt: An SG table that is already mapped to @etdev->dev, i.e. dma_map_sg*
- *       has been called.
- * @dir: DMA direction of mapping
- * @context_id: context ID for the mapping
- * @mmu_flags: the flag or'ed with EDGETPU_MMU_* macros
- *
- * Description: For chips with internal MMUs, add the required internal MMU
- * mapping for the TPU to access the DMA addresses of @sgt.
- *
- * For a chip that doesn't have an internal MMU but has the IOMMU domain AUX
- * feature, perform the necessary mapping to @context_id and return the
- * downstream DMA address.
- *
- * Caller ensures the SG table has DMA addresses as compact as possible, that is
- * if @sgt->nents is greater than 1 then the DMA addresses are not continuous.
- *
- * Returns zero on error.
- */
-tpu_addr_t edgetpu_mmu_tpu_map_sgt(struct edgetpu_dev *etdev,
-				   struct sg_table *sgt,
-				   enum dma_data_direction dir,
-				   enum edgetpu_context_id context_id,
-				   u32 mmu_flags);
-/* Unmap a TPU mapping created by edgetpu_mmu_tpu_map_sgt */
-void edgetpu_mmu_tpu_unmap_sgt(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
-			       struct sg_table *sgt,
-			       enum edgetpu_context_id context_id);
-
 /*
  * Allocates a IOMMU domain.
  *
@@ -292,4 +203,17 @@ int edgetpu_mmu_attach_domain(struct edgetpu_dev *etdev,
 void edgetpu_mmu_detach_domain(struct edgetpu_dev *etdev,
 			       struct edgetpu_iommu_domain *etdomain);
 
+/* TODO(b/281459896) Make domain comparisons internal to edgetpu-mmu.h */
+/*
+ * Returns whether mappings for a given context exist in the default domain.
+ *
+ * If a context represented by @ctx_id has been assigned the default IOMMU domain, either uniquely,
+ * or because AUX domains are not supported, this function returns true.
+ *
+ * This can be used to determine if a buffer which is already mapped for the default domain (such as
+ * coherent buffers or dma-bufs) needs to be remapped for specifically for the context.
+ */
+bool edgetpu_mmu_is_context_using_default_domain(struct edgetpu_dev *etdev,
+						 enum edgetpu_context_id ctx_id);
+
 #endif /* __EDGETPU_MMU_H__ */
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index db66dab..b2a10f9 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -16,10 +16,6 @@
 #include <gcip/gcip-pm.h>
 #include <gcip/gcip-iommu.h>
 
-#if HAS_IOVAD_BEST_FIT_ALGO
-#include <linux/dma-iommu.h>
-#endif
-
 #include "edgetpu-config.h"
 #include "edgetpu-dmabuf.h"
 #include "edgetpu-internal.h"
@@ -229,16 +225,15 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
-static void edgetpu_platform_parse_pmu(struct edgetpu_mobile_platform_dev *etmdev)
+static void edgetpu_platform_parse_pmu(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
 	struct device *dev = etdev->dev;
 	u32 reg;
 
 	if (of_find_property(dev->of_node, "pmu-status-base", NULL) &&
 	    !of_property_read_u32_index(dev->of_node, "pmu-status-base", 0, &reg)) {
-		etmdev->pmu_status = devm_ioremap(dev, reg, 0x4);
-		if (!etmdev->pmu_status)
+		etdev->pmu_status = devm_ioremap(dev, reg, 0x4);
+		if (!etdev->pmu_status)
 			etdev_err(etdev, "Using ACPM for blk status query\n");
 	} else {
 		etdev_warn(etdev, "Failed to find PMU register base\n");
@@ -354,10 +349,6 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 		goto out_cleanup_fw;
 	}
 
-#if HAS_IOVAD_BEST_FIT_ALGO
-	iommu_dma_enable_best_fit_algo(dev);
-#endif
-
 	INIT_LIST_HEAD(&etmdev->fw_ctx_list);
 	mutex_init(&etmdev->fw_ctx_list_lock);
 
@@ -365,7 +356,7 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 	 * Parses PMU before edgetpu_device_add so edgetpu_chip_pm_create can know whether to set
 	 * the is_block_down op.
 	 */
-	edgetpu_platform_parse_pmu(etmdev);
+	edgetpu_platform_parse_pmu(etdev);
 
 	ret = edgetpu_device_add(etdev, &regs, iface_params, ARRAY_SIZE(iface_params));
 	if (ret) {
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index e89e0f0..147723c 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -102,8 +102,6 @@ struct edgetpu_mobile_platform_dev {
 	int n_irq;
 	/* Array of IRQ numbers */
 	int *irq;
-	/* PMU status base address for block status, maybe NULL */
-	void __iomem *pmu_status;
 
 	/* callbacks for chip-dependent implementations */
 
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index 16e6892..506a5fd 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -236,10 +236,8 @@ void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 	int ret;
 
 	ustats = devm_kzalloc(etdev->dev, sizeof(*etdev->usage_stats), GFP_KERNEL);
-	if (!ustats) {
-		etdev_warn(etdev, "failed to allocate memory for usage stats\n");
+	if (!ustats)
 		return;
-	}
 
 	args.version = EDGETPU_USAGE_METRIC_VERSION;
 	args.dev = etdev->dev;
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c
index 75509cd..50170c6 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c
@@ -85,7 +85,11 @@ static int iovad_initialize_domain(struct gcip_iommu_domain *domain)
 	init_iova_domain(&domain->iova_space.iovad, dpool->granule,
 			 max_t(unsigned long, 1, dpool->base_daddr >> ilog2(dpool->granule)));
 
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 18, 0)
+	return iova_domain_init_rcaches(&domain->iova_space.iovad);
+#else
 	return 0;
+#endif
 }
 
 static void iovad_finalize_domain(struct gcip_iommu_domain *domain)
@@ -204,12 +208,14 @@ static ssize_t dma_iommu_map_sg(struct gcip_iommu_domain *domain, struct scatter
 	if (!nents_mapped)
 		return 0;
 
-	iova = sg_dma_address(sgl);
+	if (!domain->default_domain) {
+		iova = sg_dma_address(sgl);
 
-	ret = (ssize_t)iommu_map_sg(domain->domain, iova, sgl, nents, prot);
-	if (ret <= 0) {
-		dma_unmap_sg_attrs(domain->dev, sgl, nents, dir, attrs);
-		return 0;
+		ret = (ssize_t)iommu_map_sg(domain->domain, iova, sgl, nents, prot);
+		if (ret <= 0) {
+			dma_unmap_sg_attrs(domain->dev, sgl, nents, dir, attrs);
+			return 0;
+		}
 	}
 
 	return nents_mapped;
@@ -222,11 +228,13 @@ static void dma_iommu_unmap_sg(struct gcip_iommu_domain *domain, struct scatterl
 	size_t size = 0;
 	int i;
 
-	for_each_sg (sgl, sg, nents, i)
-		size += sg_dma_len(sg);
+	if (!domain->default_domain) {
+		for_each_sg (sgl, sg, nents, i)
+			size += sg_dma_len(sg);
 
-	if (!iommu_unmap(domain->domain, sg_dma_address(sgl), size))
-		dev_warn(domain->dev, "Failed to unmap sg");
+		if (!iommu_unmap(domain->domain, sg_dma_address(sgl), size))
+			dev_warn(domain->dev, "Failed to unmap sg");
+	}
 	dma_unmap_sg_attrs(domain->dev, sgl, nents, dir, attrs);
 }
 
@@ -443,6 +451,7 @@ struct gcip_iommu_domain *gcip_iommu_get_domain_for_dev(struct device *dev)
 
 	gdomain->dev = dev;
 	gdomain->legacy_mode = true;
+	gdomain->default_domain = true;
 
 	return gdomain;
 }
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-kci.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-kci.c
index 417b078..dd8174c 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-kci.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-kci.c
@@ -242,6 +242,13 @@ static bool gcip_kci_before_handle_resp(struct gcip_mailbox *mailbox, const void
 	return true;
 }
 
+static inline bool gcip_kci_is_block_off(struct gcip_mailbox *mailbox)
+{
+	struct gcip_kci *kci = gcip_mailbox_get_data(mailbox);
+
+	return kci->ops->is_block_off ? kci->ops->is_block_off(kci) : false;
+}
+
 static const struct gcip_mailbox_ops gcip_mailbox_ops = {
 	.get_cmd_queue_head = gcip_kci_get_cmd_queue_head,
 	.get_cmd_queue_tail = gcip_kci_get_cmd_queue_tail,
@@ -265,6 +272,7 @@ static const struct gcip_mailbox_ops gcip_mailbox_ops = {
 	.after_enqueue_cmd = gcip_kci_after_enqueue_cmd,
 	.after_fetch_resps = gcip_kci_after_fetch_resps,
 	.before_handle_resp = gcip_kci_before_handle_resp,
+	.is_block_off = gcip_kci_is_block_off,
 };
 
 /*
@@ -357,7 +365,7 @@ static int gcip_reverse_kci_remove_resp(struct gcip_reverse_kci *rkci,
 	 * Prevents the compiler from discarding and reloading its cached value additionally forces
 	 * the CPU to order against subsequent memory references.
 	 * Shamelessly stolen from:
-	 * https://www.kernel.org/doc/html/latest/core-api/circular-buffers.html
+	 * [REDACTED]
 	 */
 	head = smp_load_acquire(&rkci->head);
 	tail = rkci->tail;
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
index 4571aa9..afa67c8 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-mailbox.c
@@ -48,6 +48,8 @@
 #define RELEASE_WAIT_LIST_LOCK(irqrestore, flags)                                                  \
 	mailbox->ops->release_wait_list_lock(mailbox, irqrestore, flags)
 
+#define IS_BLOCK_OFF() (mailbox->ops->is_block_off ? mailbox->ops->is_block_off(mailbox) : false)
+
 struct gcip_mailbox_wait_list_elem {
 	struct list_head list;
 	struct gcip_mailbox_async_resp *async_resp;
@@ -331,8 +333,8 @@ static void *gcip_mailbox_fetch_responses(struct gcip_mailbox *mailbox, u32 *tot
 	void *prev_ptr = NULL; /* Temporary pointer to realloc ret. */
 	bool atomic = false;
 
-	/* Someone is working on consuming - we can leave early. */
-	if (!ACQUIRE_RESP_QUEUE_LOCK(true, &atomic))
+	/* The block is off or someone is working on consuming - we can leave early. */
+	if (IS_BLOCK_OFF() || !ACQUIRE_RESP_QUEUE_LOCK(true, &atomic))
 		goto out;
 
 	head = GET_RESP_QUEUE_HEAD();
@@ -396,7 +398,7 @@ static int gcip_mailbox_fetch_one_response(struct gcip_mailbox *mailbox, void *r
 	u32 tail;
 	bool atomic;
 
-	if (!ACQUIRE_RESP_QUEUE_LOCK(true, &atomic))
+	if (IS_BLOCK_OFF() || !ACQUIRE_RESP_QUEUE_LOCK(true, &atomic))
 		return 0;
 
 	head = GET_RESP_QUEUE_HEAD();
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c
index b9907a1..50bc74b 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-pm.c
@@ -52,6 +52,14 @@ static void gcip_pm_async_power_down_work(struct work_struct *work)
 	mutex_unlock(&pm->lock);
 }
 
+/* Worker for async gcip_pm_put(). */
+static void gcip_pm_async_put_work(struct work_struct *work)
+{
+	struct gcip_pm *pm = container_of(work, struct gcip_pm, put_async_work);
+
+	gcip_pm_put(pm);
+}
+
 struct gcip_pm *gcip_pm_create(const struct gcip_pm_args *args)
 {
 	struct gcip_pm *pm;
@@ -73,6 +81,7 @@ struct gcip_pm *gcip_pm_create(const struct gcip_pm_args *args)
 
 	mutex_init(&pm->lock);
 	INIT_DELAYED_WORK(&pm->power_down_work, gcip_pm_async_power_down_work);
+	INIT_WORK(&pm->put_async_work, gcip_pm_async_put_work);
 
 	if (pm->after_create) {
 		ret = pm->after_create(pm->data);
@@ -186,6 +195,16 @@ unlock:
 	mutex_unlock(&pm->lock);
 }
 
+void gcip_pm_put_async(struct gcip_pm *pm)
+{
+	schedule_work(&pm->put_async_work);
+}
+
+void gcip_pm_flush_put_work(struct gcip_pm *pm)
+{
+	flush_work(&pm->put_async_work);
+}
+
 int gcip_pm_get_count(struct gcip_pm *pm)
 {
 	if (!pm)
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h
index 1797f94..34f6efa 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h
@@ -104,6 +104,7 @@ struct gcip_iommu_domain {
 	struct gcip_iommu_domain_pool *domain_pool;
 	struct iommu_domain *domain;
 	bool legacy_mode;
+	bool default_domain;
 	union {
 		struct iova_domain iovad;
 		struct gcip_mem_pool mem_pool;
@@ -238,13 +239,13 @@ static inline bool gcip_iommu_domain_is_legacy_mode(struct gcip_iommu_domain *do
  *               00 = DMA_BIDIRECTIONAL (host/device can write buffer)
  *               01 = DMA_TO_DEVICE     (host can write buffer)
  *               10 = DMA_FROM_DEVICE   (device can write buffer)
- *               (See https://docs.kernel.org/core-api/dma-api-howto.html#dma-direction)
+ *               (See [REDACTED]
  *   [2:2]   - Coherent Mapping:
  *               0 = Create non-coherent mappings of the buffer.
  *               1 = Create coherent mappings of the buffer.
  *   [12:3]  - DMA_ATTR:
  *               Not used in the non-legacy mode.
- *               (See https://www.kernel.org/doc/Documentation/core-api/dma-attributes.rst)
+ *               (See [REDACTED]
  *   [63:13] - RESERVED
  *               Set RESERVED bits to 0 to ensure backwards compatibility.
  *
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
index 1cfc82e..74670de 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
@@ -233,6 +233,11 @@ struct gcip_kci_ops {
 	 * Context: normal.
 	 */
 	int (*update_usage)(struct gcip_kci *kci);
+	/*
+	 * Checks if the block is off.
+	 * Context: in_interrupt().
+	 */
+	bool (*is_block_off)(struct gcip_kci *kci);
 };
 
 struct gcip_kci {
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
index af48ba6..b16c15f 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
@@ -344,6 +344,11 @@ struct gcip_mailbox_ops {
 	 * Context: normal and in_interrupt().
 	 */
 	void (*release_awaiter_data)(void *data);
+	/*
+	 * Checks if the block is off.
+	 * Context: in_interrupt()
+	 */
+	bool (*is_block_off)(struct gcip_mailbox *mailbox);
 };
 
 struct gcip_mailbox {
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
index 1e6ce05..7e3a3e4 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
@@ -23,6 +23,8 @@ struct gcip_pm {
 	int count;
 	/* Flag indicating a deferred power down is pending. Protected by @lock */
 	bool power_down_pending;
+	/* The worker to asynchronously call gcip_pm_put(). */
+	struct work_struct put_async_work;
 
 	/* Callbacks. See struct gcip_pm_args. */
 	void *data;
@@ -99,6 +101,12 @@ int gcip_pm_get(struct gcip_pm *pm);
  */
 void gcip_pm_put(struct gcip_pm *pm);
 
+/* Schedules an asynchronous job to execute gcip_pm_put(). */
+void gcip_pm_put_async(struct gcip_pm *pm);
+
+/* Flushes the pending pm_put work if any. */
+void gcip_pm_flush_put_work(struct gcip_pm *pm);
+
 /* Gets the power up counter. Note that this is checked without PM lock. */
 int gcip_pm_get_count(struct gcip_pm *pm);
 
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index b903ba1..d247c06 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -323,7 +323,7 @@ long edgetpu_soc_pm_get_rate(struct edgetpu_dev *etdev, int flags)
 
 	/*
 	 * Below values must match the CMU PLL (pll_con3_pll_tpu) values in the spec and firmware.
-	 * See https://drive.google.com/file/d/16S9yxmGwkOltdO2w4dC8tpAt99chn-aq/view and
+	 * See [REDACTED] and
 	 * power_manager.cc for more details.
 	 */
 	switch (TO_PLL_DIV_M(pll_con3)) {
diff --git a/drivers/edgetpu/rio-pm.c b/drivers/edgetpu/rio-pm.c
index ac771ec..c50949a 100644
--- a/drivers/edgetpu/rio-pm.c
+++ b/drivers/edgetpu/rio-pm.c
@@ -144,14 +144,13 @@ static int rio_lpm_up(struct edgetpu_dev *etdev)
 
 static bool rio_is_block_down(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	int timeout_cnt = 0;
 	int curr_state;
 
 	do {
 		/* Delay 20us per retry till blk shutdown finished */
 		usleep_range(SHUTDOWN_DELAY_US_MIN, SHUTDOWN_DELAY_US_MAX);
-		curr_state = readl(etmdev->pmu_status);
+		curr_state = readl(etdev->pmu_status);
 		if (!curr_state)
 			return true;
 		timeout_cnt++;
@@ -173,7 +172,7 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 
 	platform_pwr->lpm_up = rio_lpm_up;
 	platform_pwr->lpm_down = rio_lpm_down;
-	if (etmdev->pmu_status)
+	if (etdev->pmu_status)
 		platform_pwr->is_block_down = rio_is_block_down;
 	platform_pwr->post_fw_start = rio_post_fw_start;