Merge android-gs-pixel-mainline into android-gs-pixel-5.10

Signed-off-by: Petri Gynther <pgynther@google.com>
Change-Id: I976fd588cd8c05b6760c4011f5e8da750f466f90

14 files changed, 957 insertions, 834 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 5361fa3..f6c7d96 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -12,15 +12,15 @@ endif
 
 edgetpu-fw-objs := edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-shared-fw.o
 edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o $(edgetpu-fw-objs)
-abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-iommu.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o abrolhos-debug-dump.o abrolhos-usage-stats.o $(edgetpu-objs)
+abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o abrolhos-iommu.o abrolhos-debug-dump.o abrolhos-usage-stats.o $(edgetpu-objs)
 CFLAGS_abrolhos-fs.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-core.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-device-group.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-firmware.o := -DCONFIG_ABROLHOS=1
+CFLAGS_abrolhos-iommu.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-platform.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-pm.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-thermal.o := -DCONFIG_ABROLHOS=1
-CFLAGS_abrolhos-iommu.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-debug-dump.o := -DCONFIG_ABROLHOS=1
 CFLAGS_abrolhos-usage-stats.o := -DCONFIG_ABROLHOS=1
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index 0391dcf..fd195b6 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -16,7 +16,9 @@ endif
 edgetpu-fw-objs := edgetpu-firmware-util.o edgetpu-shared-fw.o edgetpu-firmware.o
 edgetpu-objs	:= edgetpu-core.o edgetpu-mailbox.o edgetpu-kci.o edgetpu-device-group.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o $(edgetpu-fw-objs)
 
-abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-fs.o abrolhos-platform.o abrolhos-iommu.o abrolhos-thermal.o abrolhos-pm.o abrolhos-debug-dump.o abrolhos-usage-stats.o $(edgetpu-objs)
+edgetpu-mobile-objs := edgetpu-google-iommu.o $(edgetpu-objs)
+
+abrolhos-objs	:= abrolhos-device.o abrolhos-firmware.o edgetpu-fs.o abrolhos-platform.o abrolhos-thermal.o abrolhos-pm.o abrolhos-debug-dump.o abrolhos-usage-stats.o $(edgetpu-mobile-objs)
 
 KBUILD_OPTIONS += CONFIG_ABROLHOS=m
 
diff --git a/drivers/edgetpu/abrolhos-iommu.c b/drivers/edgetpu/abrolhos-iommu.c
index 947e7c4..56a9f6d 100644
--- a/drivers/edgetpu/abrolhos-iommu.c
+++ b/drivers/edgetpu/abrolhos-iommu.c
@@ -1,764 +1,2 @@
 // SPDX-License-Identifier: GPL-2.0
-/*
- * Edge TPU IOMMU interface.
- *
- * Copyright (C) 2019 Google, Inc.
- */
-
-#include <linux/device.h>
-#include <linux/dma-mapping.h>
-#include <linux/idr.h>
-#include <linux/iommu.h>
-#include <linux/scatterlist.h>
-#include <linux/slab.h>
-#include <linux/types.h>
-#include <linux/version.h>
-
-#include "abrolhos-platform.h"
-#include "edgetpu-internal.h"
-#include "edgetpu-mapping.h"
-#include "edgetpu-mmu.h"
-
-struct edgetpu_iommu {
-	struct iommu_group *iommu_group;
-	/*
-	 * IOMMU domains currently attached.
-	 * NULL for a slot that doesn't have an attached domain.
-	 */
-	struct iommu_domain *domains[EDGETPU_NCONTEXTS];
-	/*
-	 * Records all domains currently allocated, to support IOMMU (un)mapping
-	 * when the domain is not attached.
-	 */
-	struct idr domain_pool;
-	struct mutex pool_lock;		/* protects access of @domain_pool */
-	bool context_0_default;		/* is context 0 domain the default? */
-	bool aux_enabled;
-};
-
-struct edgetpu_iommu_map_params {
-	int prot;
-	size_t size;
-	struct iommu_domain *domain;
-};
-
-/*
- * Return context ID enumeration value as a Process Address Space ID.
- * Caller ensures context_id is valid, i.e. does not equal to
- * EDGETPU_CONTEXT_INVALID or OR'ed with EDGETPU_CONTEXT_DOMAIN_TOKEN.
- */
-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)
-{
-	struct iommu_domain *domain;
-
-	mutex_lock(&etiommu->pool_lock);
-	domain = idr_find(&etiommu->domain_pool, token);
-	mutex_unlock(&etiommu->pool_lock);
-	return domain;
-}
-
-static struct 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 edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	uint pasid;
-
-	if (ctx_id == EDGETPU_CONTEXT_INVALID)
-		return NULL;
-	if (ctx_id & EDGETPU_CONTEXT_DOMAIN_TOKEN)
-		return get_domain_by_token(
-			etiommu, ctx_id ^ EDGETPU_CONTEXT_DOMAIN_TOKEN);
-	pasid = context_id_to_pasid(ctx_id);
-	if (pasid < EDGETPU_NCONTEXTS)
-		domain = etiommu->domains[pasid];
-
-	/* Fall back to default domain. */
-	if (!domain)
-		domain = iommu_get_domain_for_dev(dev);
-	return domain;
-}
-
-/*
- * Kernel 5.3 introduced iommu_register_device_fault_handler
- */
-
-#if KERNEL_VERSION(5, 3, 0) <=  LINUX_VERSION_CODE
-
-static int edgetpu_iommu_dev_fault_handler(struct iommu_fault *fault,
-					   void *token)
-{
-	struct edgetpu_dev *etdev = (struct edgetpu_dev *)token;
-
-	if (fault->type == IOMMU_FAULT_DMA_UNRECOV) {
-		etdev_err(etdev, "Unrecoverable IOMMU fault!\n");
-		etdev_err(etdev, "Reason = %08X\n", fault->event.reason);
-		etdev_err(etdev, "flags = %08X\n", fault->event.flags);
-		etdev_err(etdev, "pasid = %08X\n", fault->event.pasid);
-		etdev_err(etdev, "perms = %08X\n", fault->event.perm);
-		etdev_err(etdev, "addr = %llX\n", fault->event.addr);
-		etdev_err(etdev, "fetch_addr = %llX\n",
-			  fault->event.fetch_addr);
-	} else if (fault->type == IOMMU_FAULT_PAGE_REQ) {
-		etdev_err(etdev, "IOMMU page request fault!\n");
-		etdev_err(etdev, "flags = %08X\n", fault->prm.flags);
-		etdev_err(etdev, "pasid = %08X\n", fault->prm.pasid);
-		etdev_err(etdev, "grpid = %08X\n", fault->prm.grpid);
-		etdev_err(etdev, "perms = %08X\n", fault->prm.perm);
-		etdev_err(etdev, "addr = %llX\n", fault->prm.addr);
-	}
-	// Tell the IOMMU driver to carry on
-	return -EAGAIN;
-}
-
-static int
-edgetpu_register_iommu_device_fault_handler(struct edgetpu_dev *etdev)
-{
-	etdev_dbg(etdev, "Registering IOMMU device fault handler\n");
-	return iommu_register_device_fault_handler(
-		etdev->dev, edgetpu_iommu_dev_fault_handler, etdev);
-}
-
-static int
-edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etdev)
-{
-	etdev_dbg(etdev, "Unregistering IOMMU device fault handler\n");
-	return iommu_unregister_device_fault_handler(etdev->dev);
-}
-
-#else /* kernel version before 5.3 */
-
-static int
-edgetpu_register_iommu_device_fault_handler(struct edgetpu_dev *etdev)
-{
-	return 0;
-}
-
-static int
-edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etdev)
-{
-	return 0;
-}
-
-#endif /* KERNEL_VERSION(5, 3, 0) <=  LINUX_VERSION_CODE */
-
-/* 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;
-
-	iommu_domain_free(domain);
-	return 0;
-}
-
-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;
-
-	dev_err(dev, "IOMMU fault on address %08lX. PASID = %u flags = %08X",
-		iova, etdomain->pasid, flags);
-	// Tell the IOMMU driver we are OK with this fault
-	return 0;
-}
-
-static void edgetpu_init_etdomain(struct edgetpu_iommu_domain *etdomain,
-				  struct iommu_domain *domain,
-				  int token)
-{
-	etdomain->iommu_domain = domain;
-	etdomain->pasid = IOMMU_PASID_INVALID;
-	etdomain->token = token;
-	iommu_set_fault_handler(domain, edgetpu_iommu_fault_handler, etdomain);
-}
-
-/*
- * Expect a default domain was already allocated for the group. If not try to
- * use the domain AUX feature to allocate one.
- */
-static int check_default_domain(struct edgetpu_dev *etdev,
-				struct edgetpu_iommu *etiommu)
-{
-	struct iommu_domain *domain;
-	int ret;
-	uint pasid;
-
-	domain = iommu_get_domain_for_dev(etdev->dev);
-	/* if default domain exists then we are done */
-	if (domain) {
-		etiommu->context_0_default = true;
-		goto out;
-	}
-	etdev_warn(etdev, "device group has no default iommu domain\n");
-	/* no default domain and no AUX - we can't have any domain */
-	if (!etiommu->aux_enabled)
-		return -EINVAL;
-
-	domain = iommu_domain_alloc(etdev->dev->bus);
-	if (!domain) {
-		etdev_warn(etdev, "iommu domain alloc failed");
-		return -EINVAL;
-	}
-	ret = iommu_aux_attach_device(domain, etdev->dev);
-	if (ret) {
-		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
-		iommu_domain_free(domain);
-		return ret;
-	}
-	pasid = iommu_aux_get_pasid(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);
-		iommu_domain_free(domain);
-		return -EINVAL;
-	}
-out:
-	etiommu->domains[0] = domain;
-	return 0;
-}
-
-/* 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)
-{
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
-	struct edgetpu_iommu *etiommu;
-	int ret;
-
-	etiommu = kzalloc(sizeof(*etiommu), GFP_KERNEL);
-	if (!etiommu)
-		return -ENOMEM;
-	idr_init(&etiommu->domain_pool);
-	mutex_init(&etiommu->pool_lock);
-	etiommu->iommu_group = iommu_group_get(etdev->dev);
-	if (etiommu->iommu_group) {
-		iommu_group_set_name(etiommu->iommu_group, "edgetpu");
-		dev_dbg(etdev->dev, "iommu group id %d setup\n",
-			iommu_group_id(etiommu->iommu_group));
-	} else {
-		dev_warn(etdev->dev, "device has no iommu group\n");
-	}
-
-	iommu_dev_enable_feature(etdev->dev, IOMMU_DEV_FEAT_AUX);
-	if (!iommu_dev_feature_enabled(etdev->dev, IOMMU_DEV_FEAT_AUX))
-		etdev_warn(etdev, "AUX domains not supported\n");
-	else
-		etiommu->aux_enabled = true;
-	ret = check_default_domain(etdev, etiommu);
-	if (ret)
-		goto err_free;
-
-	ret = edgetpu_register_iommu_device_fault_handler(etdev);
-	if (ret)
-		etdev_warn(etdev, "Failed to register fault handler! (%d)\n",
-			   ret);
-
-	/* etiommu initialization done */
-	etdev->mmu_cookie = etiommu;
-	if (!edgetpu_pdev->csr_iova)
-		goto success;
-
-	etdev_dbg(etdev, "Mapping device CSRs: %llX -> %llX (%lu bytes)\n",
-		  edgetpu_pdev->csr_iova, edgetpu_pdev->csr_paddr,
-		  edgetpu_pdev->csr_size);
-
-	/* Add an IOMMU translation for the CSR region */
-	ret = edgetpu_mmu_add_translation(etdev, edgetpu_pdev->csr_iova,
-					  edgetpu_pdev->csr_paddr,
-					  edgetpu_pdev->csr_size,
-					  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:
-	return 0;
-
-err_free:
-	kfree(etiommu);
-	return ret;
-}
-
-void edgetpu_mmu_reset(struct edgetpu_dev *etdev)
-{
-	/* If need to reset IOMMU driver can issue here. */
-}
-
-void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
-{
-	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	int i, ret;
-
-	if (!etiommu)
-		return;
-
-	if (edgetpu_pdev->csr_iova) {
-		edgetpu_mmu_remove_translation(&edgetpu_pdev->edgetpu_dev,
-					       edgetpu_pdev->csr_iova,
-					       edgetpu_pdev->csr_size,
-					       EDGETPU_CONTEXT_KCI);
-	}
-	edgetpu_pdev->csr_iova = 0;
-
-	ret = edgetpu_unregister_iommu_device_fault_handler(etdev);
-	if (ret)
-		etdev_warn(etdev,
-			   "Failed to unregister device fault handler (%d)\n",
-			   ret);
-	edgetpu_mmu_reset(etdev);
-
-	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->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])
-		iommu_domain_free(etiommu->domains[0]);
-
-	idr_for_each(&etiommu->domain_pool, edgetpu_idr_free_domain_callback,
-		     NULL);
-	idr_destroy(&etiommu->domain_pool);
-	kfree(etiommu);
-	etdev->mmu_cookie = NULL;
-}
-
-int edgetpu_mmu_reattach(struct edgetpu_dev *etdev)
-{
-	return 0;
-}
-
-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)
-{
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	size_t size = 0;
-	int prot = __dma_dir_to_iommu_prot(map->dir);
-	struct iommu_domain *domain;
-	int i;
-	struct scatterlist *sg;
-
-	if (!etiommu)
-		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;
-	return 0;
-}
-
-int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
-		    enum edgetpu_context_id context_id, u32 mmu_flags)
-{
-	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);
-
-	if (ret)
-		return ret;
-
-	if (mmu_flags & EDGETPU_MMU_64)
-		dev_warn_once(etdev->dev,
-			      "%s: 64-bit addressing is not supported",
-			      __func__);
-
-	ret = dma_map_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.nents,
-			       edgetpu_host_dma_dir(map->dir), map->dma_attrs);
-	if (!ret)
-		return -EINVAL;
-	map->sgt.nents = ret;
-	iova = sg_dma_address(map->sgt.sgl);
-
-	/*
-	 * 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) {
-		if (!iommu_map_sg(params.domain, iova, map->sgt.sgl,
-				  map->sgt.orig_nents, params.prot)) {
-			/* Undo the mapping in the default domain */
-			dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl,
-					   map->sgt.orig_nents,
-					   edgetpu_host_dma_dir(map->dir),
-					   DMA_ATTR_SKIP_CPU_SYNC);
-			return -ENOMEM;
-		}
-	}
-
-	map->device_address = iova;
-	return 0;
-}
-
-void edgetpu_mmu_unmap(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
-		       enum edgetpu_context_id context_id)
-{
-	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);
-	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);
-	}
-
-	/* Undo the mapping in the default domain */
-	dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.orig_nents,
-			   edgetpu_host_dma_dir(map->dir), map->dma_attrs);
-}
-
-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 tpu_addr_t orig_iova = iova;
-	struct scatterlist *sg;
-	int i;
-	int ret;
-
-	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
-		ret = edgetpu_mmu_add_translation(etdev, iova, sg_phys(sg),
-						  sg->length, prot, context_id);
-		if (ret)
-			goto error;
-		iova += sg->length;
-	}
-	return 0;
-
-error:
-	edgetpu_mmu_remove_translation(etdev, orig_iova, iova - orig_iova,
-				       context_id);
-	return ret;
-}
-
-void edgetpu_mmu_unmap_iova_sgt_attrs(struct edgetpu_dev *etdev,
-				      tpu_addr_t iova, struct sg_table *sgt,
-				      enum dma_data_direction dir,
-				      enum edgetpu_context_id context_id,
-				      unsigned long attrs)
-{
-	size_t size = 0;
-	struct scatterlist *sg;
-	int i;
-
-	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i)
-		size += sg->length;
-	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;
-
-	domain = get_domain_by_context_id(etdev, context_id);
-	if (!domain)
-		return -ENODEV;
-	return iommu_map(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;
-
-	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 = __dma_dir_to_iommu_prot(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 = __dma_dir_to_iommu_prot(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));
-}
-
-void edgetpu_mmu_use_dev_dram(struct edgetpu_dev *etdev, bool use_dev_dram)
-{
-}
-
-/* to be returned when domain aux is not supported */
-static struct edgetpu_iommu_domain invalid_etdomain = {
-	.pasid = IOMMU_PASID_INVALID,
-	.token = EDGETPU_DOMAIN_TOKEN_END,
-};
-
-struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
-{
-	struct edgetpu_iommu_domain *etdomain =
-		kzalloc(sizeof(*etdomain), GFP_KERNEL);
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	struct iommu_domain *domain;
-	int token;
-
-	if (!etdomain)
-		return NULL;
-	if (!etiommu->aux_enabled)
-		return &invalid_etdomain;
-	domain = iommu_domain_alloc(etdev->dev->bus);
-	if (!domain) {
-		etdev_warn(etdev, "iommu domain alloc failed");
-		return NULL;
-	}
-
-	mutex_lock(&etiommu->pool_lock);
-	token = idr_alloc(&etiommu->domain_pool, domain, 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);
-		iommu_domain_free(domain);
-		return NULL;
-	}
-	edgetpu_init_etdomain(etdomain, domain, token);
-	return etdomain;
-}
-
-void edgetpu_mmu_free_domain(struct edgetpu_dev *etdev,
-			     struct edgetpu_iommu_domain *etdomain)
-{
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-
-	if (!etdomain || etdomain == &invalid_etdomain)
-		return;
-	if (etdomain->pasid != IOMMU_PASID_INVALID) {
-		etdev_warn(etdev, "Domain should be detached before free");
-		edgetpu_mmu_detach_domain(etdev, etdomain);
-	}
-	mutex_lock(&etiommu->pool_lock);
-	idr_remove(&etiommu->domain_pool, etdomain->token);
-	mutex_unlock(&etiommu->pool_lock);
-	iommu_domain_free(etdomain->iommu_domain);
-	kfree(etdomain);
-}
-
-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;
-	int ret;
-	uint pasid;
-
-	/* Changes nothing if domain AUX is not supported. */
-	if (!etiommu->aux_enabled)
-		return 0;
-	if (etdomain->pasid != IOMMU_PASID_INVALID)
-		return -EINVAL;
-	domain = etdomain->iommu_domain;
-	ret = iommu_aux_attach_device(domain, etdev->dev);
-	if (ret) {
-		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
-		return ret;
-	}
-	pasid = iommu_aux_get_pasid(domain, etdev->dev);
-	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS) {
-		etdev_warn(etdev, "Invalid PASID %d returned from iommu",
-			   pasid);
-		ret = -EINVAL;
-		goto err_detach;
-	}
-	/* the IOMMU driver returned a duplicate PASID */
-	if (etiommu->domains[pasid]) {
-		ret = -EBUSY;
-		goto err_detach;
-	}
-	etiommu->domains[pasid] = domain;
-	etdomain->pasid = pasid;
-	return 0;
-err_detach:
-	iommu_aux_detach_device(domain, etdev->dev);
-	return ret;
-}
-
-void edgetpu_mmu_detach_domain(struct edgetpu_dev *etdev,
-			       struct edgetpu_iommu_domain *etdomain)
-{
-	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
-	uint pasid = etdomain->pasid;
-
-	if (!etiommu->aux_enabled)
-		return;
-	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS)
-		return;
-	etiommu->domains[pasid] = NULL;
-	etdomain->pasid = IOMMU_PASID_INVALID;
-	iommu_aux_detach_device(etdomain->iommu_domain, etdev->dev);
-}
+#include "edgetpu-google-iommu.c"
\ No newline at end of file
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 04b42b5..12bbe3f 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -434,7 +434,8 @@ static int abrolhos_power_up(struct edgetpu_pm *etpm)
 }
 
 static void
-abrolhos_pm_shutdown_firmware(struct edgetpu_dev *etdev,
+abrolhos_pm_shutdown_firmware(struct abrolhos_platform_dev *etpdev,
+			      struct edgetpu_dev *etdev,
 			      struct abrolhos_platform_dev *edgetpu_pdev)
 {
 	if (!edgetpu_pchannel_power_down(etdev, false))
@@ -449,6 +450,8 @@ abrolhos_pm_shutdown_firmware(struct edgetpu_dev *etdev,
 
 	cancel_work_sync(&etdev->kci->work);
 	etdev_warn(etdev, "Forcing shutdown through power policy\n");
+	/* Request GSA shutdown to make sure the R52 core is reset */
+	gsa_send_tpu_cmd(etpdev->gsa_dev, GSA_TPU_SHUTDOWN);
 	abrolhos_pwr_policy_set(edgetpu_pdev, TPU_OFF);
 	pm_runtime_put_sync(etdev->dev);
 	/*
@@ -483,7 +486,8 @@ static void abrolhos_power_down(struct edgetpu_pm *etpm)
 	if (etdev->kci && edgetpu_firmware_status_locked(etdev) == FW_VALID) {
 		/* Update usage stats before we power off fw. */
 		edgetpu_kci_update_usage(etdev);
-		abrolhos_pm_shutdown_firmware(etdev, edgetpu_pdev);
+		abrolhos_pm_shutdown_firmware(edgetpu_pdev, etdev,
+					      edgetpu_pdev);
 		cancel_work_sync(&etdev->kci->work);
 	}
 
diff --git a/drivers/edgetpu/abrolhos/config-tpu-cpu.h b/drivers/edgetpu/abrolhos/config-tpu-cpu.h
index 81e6187..123dc57 100644
--- a/drivers/edgetpu/abrolhos/config-tpu-cpu.h
+++ b/drivers/edgetpu/abrolhos/config-tpu-cpu.h
@@ -20,7 +20,8 @@
 
 /* Power Control signals for P-channel interface. */
 #define EDGETPU_REG_POWER_CONTROL			0xA0008
-#define PSTATE						(1 << 0)
+#define PSTATE_SHIFT					0
+#define PSTATE						(1 << PSTATE_SHIFT)
 #define PREQ						(1 << 1)
 #define PDENY						(1 << 2)
 #define PACCEPT						(1 << 3)
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 9735b28..a5c05e1 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -33,6 +33,10 @@
 #include "edgetpu-usage-stats.h"
 #include "edgetpu.h"
 
+#define UNLOCK(client) mutex_unlock(&client->group_lock)
+#define LOCK_IN_GROUP(client)                           \
+	({ mutex_lock(&client->group_lock); client->group ? 0 : -EINVAL; })
+
 static atomic_t single_dev_count = ATOMIC_INIT(-1);
 
 /* TODO(b/156444816): Check permission. */
@@ -98,7 +102,7 @@ static const struct vm_operations_struct edgetpu_vma_ops = {
 /* Map exported device CSRs or queue into user space. */
 int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 {
-	int ret;
+	int ret = 0;
 
 	if (vma->vm_start & ~PAGE_MASK) {
 		etdev_dbg(client->etdev,
@@ -137,19 +141,17 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 		return edgetpu_mmap_telemetry_buffer(
 			client->etdev, EDGETPU_TELEMETRY_TRACE, vma);
 
-	mutex_lock(&client->group_lock);
-	if (!client->group) {
-		mutex_unlock(&client->group_lock);
-		return -EINVAL;
-	}
-
 	switch (vma->vm_pgoff) {
 	case EDGETPU_MMAP_CSR_OFFSET >> PAGE_SHIFT:
 		mutex_lock(&client->wakelock.lock);
-		if (!client->wakelock.req_count)
+		if (!client->wakelock.req_count) {
 			ret = -EAGAIN;
-		else
-			ret = edgetpu_mmap_csr(client->group, vma);
+		} else {
+			ret = LOCK_IN_GROUP(client);
+			if (!ret)
+				ret = edgetpu_mmap_csr(client->group, vma);
+			UNLOCK(client);
+		}
 		if (!ret)
 			client->wakelock.csr_map_count++;
 		etdev_dbg(client->etdev, "%s: mmap CSRS. count = %u ret = %d\n",
@@ -157,17 +159,23 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 		mutex_unlock(&client->wakelock.lock);
 		break;
 	case EDGETPU_MMAP_CMD_QUEUE_OFFSET >> PAGE_SHIFT:
-		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma);
+		ret = LOCK_IN_GROUP(client);
+		if (!ret)
+			ret = edgetpu_mmap_queue(client->group,
+						 MAILBOX_CMD_QUEUE, vma);
+		UNLOCK(client);
 		break;
 	case EDGETPU_MMAP_RESP_QUEUE_OFFSET >> PAGE_SHIFT:
-		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE,
-					 vma);
+		ret = LOCK_IN_GROUP(client);
+		if (!ret)
+			ret = edgetpu_mmap_queue(client->group,
+						 MAILBOX_RESP_QUEUE, vma);
+		UNLOCK(client);
 		break;
 	default:
 		ret = -EINVAL;
 		break;
 	}
-	mutex_unlock(&client->group_lock);
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index cc6aa14..5bf795d 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -643,7 +643,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 = arg->size;
+	const u64 size = PAGE_ALIGN(arg->size);
 	const enum dma_data_direction dir =
 		edgetpu_host_dma_dir(flags & EDGETPU_MAP_DIR_MASK);
 	struct edgetpu_dev *etdev;
@@ -795,14 +795,14 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 			bmap->dmabufs[i] = dmabuf;
 		}
 	}
-	bmap->size = arg->size;
+	bmap->size = PAGE_ALIGN(arg->size);
 	for (i = 0; i < group->n_clients; i++) {
 		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,
-						   arg->size, dir);
+						   bmap->size, dir);
 		if (ret)
 			goto err_release_bmap;
 	}
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
new file mode 100644
index 0000000..b62ecea
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -0,0 +1,774 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Edge TPU IOMMU interface.
+ *
+ * Copyright (C) 2019 Google, Inc.
+ */
+
+#include <linux/device.h>
+#include <linux/dma-mapping.h>
+#include <linux/idr.h>
+#include <linux/iommu.h>
+#include <linux/scatterlist.h>
+#include <linux/slab.h>
+#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"
+
+struct edgetpu_iommu {
+	struct iommu_group *iommu_group;
+	/*
+	 * IOMMU domains currently attached.
+	 * NULL for a slot that doesn't have an attached domain.
+	 */
+	struct iommu_domain *domains[EDGETPU_NCONTEXTS];
+	/*
+	 * Records all domains currently allocated, to support IOMMU (un)mapping
+	 * when the domain is not attached.
+	 */
+	struct idr domain_pool;
+	struct mutex pool_lock;		/* protects access of @domain_pool */
+	bool context_0_default;		/* is context 0 domain the default? */
+	bool aux_enabled;
+};
+
+struct edgetpu_iommu_map_params {
+	int prot;
+	size_t size;
+	struct iommu_domain *domain;
+};
+
+/*
+ * Return context ID enumeration value as a Process Address Space ID.
+ * Caller ensures context_id is valid, i.e. does not equal to
+ * EDGETPU_CONTEXT_INVALID or OR'ed with EDGETPU_CONTEXT_DOMAIN_TOKEN.
+ */
+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)
+{
+	struct iommu_domain *domain;
+
+	mutex_lock(&etiommu->pool_lock);
+	domain = idr_find(&etiommu->domain_pool, token);
+	mutex_unlock(&etiommu->pool_lock);
+	return domain;
+}
+
+static struct 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 edgetpu_iommu *etiommu = etdev->mmu_cookie;
+	uint pasid;
+
+	if (ctx_id == EDGETPU_CONTEXT_INVALID)
+		return NULL;
+	if (ctx_id & EDGETPU_CONTEXT_DOMAIN_TOKEN)
+		return get_domain_by_token(
+			etiommu, ctx_id ^ EDGETPU_CONTEXT_DOMAIN_TOKEN);
+	pasid = context_id_to_pasid(ctx_id);
+	if (pasid < EDGETPU_NCONTEXTS)
+		domain = etiommu->domains[pasid];
+
+	/* Fall back to default domain. */
+	if (!domain)
+		domain = iommu_get_domain_for_dev(dev);
+	return domain;
+}
+
+/*
+ * Kernel 5.3 introduced iommu_register_device_fault_handler
+ */
+
+#if KERNEL_VERSION(5, 3, 0) <=  LINUX_VERSION_CODE
+
+static int edgetpu_iommu_dev_fault_handler(struct iommu_fault *fault,
+					   void *token)
+{
+	struct edgetpu_dev *etdev = (struct edgetpu_dev *)token;
+
+	if (fault->type == IOMMU_FAULT_DMA_UNRECOV) {
+		etdev_err(etdev, "Unrecoverable IOMMU fault!\n");
+		etdev_err(etdev, "Reason = %08X\n", fault->event.reason);
+		etdev_err(etdev, "flags = %08X\n", fault->event.flags);
+		etdev_err(etdev, "pasid = %08X\n", fault->event.pasid);
+		etdev_err(etdev, "perms = %08X\n", fault->event.perm);
+		etdev_err(etdev, "addr = %llX\n", fault->event.addr);
+		etdev_err(etdev, "fetch_addr = %llX\n",
+			  fault->event.fetch_addr);
+	} else if (fault->type == IOMMU_FAULT_PAGE_REQ) {
+		etdev_err(etdev, "IOMMU page request fault!\n");
+		etdev_err(etdev, "flags = %08X\n", fault->prm.flags);
+		etdev_err(etdev, "pasid = %08X\n", fault->prm.pasid);
+		etdev_err(etdev, "grpid = %08X\n", fault->prm.grpid);
+		etdev_err(etdev, "perms = %08X\n", fault->prm.perm);
+		etdev_err(etdev, "addr = %llX\n", fault->prm.addr);
+	}
+	// Tell the IOMMU driver to carry on
+	return -EAGAIN;
+}
+
+static int
+edgetpu_register_iommu_device_fault_handler(struct edgetpu_dev *etdev)
+{
+	etdev_dbg(etdev, "Registering IOMMU device fault handler\n");
+	return iommu_register_device_fault_handler(
+		etdev->dev, edgetpu_iommu_dev_fault_handler, etdev);
+}
+
+static int
+edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etdev)
+{
+	etdev_dbg(etdev, "Unregistering IOMMU device fault handler\n");
+	return iommu_unregister_device_fault_handler(etdev->dev);
+}
+
+#else /* kernel version before 5.3 */
+
+static int
+edgetpu_register_iommu_device_fault_handler(struct edgetpu_dev *etdev)
+{
+	return 0;
+}
+
+static int
+edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etdev)
+{
+	return 0;
+}
+
+#endif /* KERNEL_VERSION(5, 3, 0) <=  LINUX_VERSION_CODE */
+
+/* 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;
+
+	iommu_domain_free(domain);
+	return 0;
+}
+
+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;
+
+	dev_err(dev, "IOMMU fault on address %08lX. PASID = %u flags = %08X",
+		iova, etdomain->pasid, flags);
+	// Tell the IOMMU driver we are OK with this fault
+	return 0;
+}
+
+static void edgetpu_init_etdomain(struct edgetpu_iommu_domain *etdomain,
+				  struct iommu_domain *domain,
+				  int token)
+{
+	etdomain->iommu_domain = domain;
+	etdomain->pasid = IOMMU_PASID_INVALID;
+	etdomain->token = token;
+	iommu_set_fault_handler(domain, edgetpu_iommu_fault_handler, etdomain);
+}
+
+/*
+ * Expect a default domain was already allocated for the group. If not try to
+ * use the domain AUX feature to allocate one.
+ */
+static int check_default_domain(struct edgetpu_dev *etdev,
+				struct edgetpu_iommu *etiommu)
+{
+	struct iommu_domain *domain;
+	int ret;
+	uint pasid;
+
+	domain = iommu_get_domain_for_dev(etdev->dev);
+	/* if default domain exists then we are done */
+	if (domain) {
+		etiommu->context_0_default = true;
+		goto out;
+	}
+	etdev_warn(etdev, "device group has no default iommu domain\n");
+	/* no default domain and no AUX - we can't have any domain */
+	if (!etiommu->aux_enabled)
+		return -EINVAL;
+
+	domain = iommu_domain_alloc(etdev->dev->bus);
+	if (!domain) {
+		etdev_warn(etdev, "iommu domain alloc failed");
+		return -EINVAL;
+	}
+	ret = iommu_aux_attach_device(domain, etdev->dev);
+	if (ret) {
+		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
+		iommu_domain_free(domain);
+		return ret;
+	}
+	pasid = iommu_aux_get_pasid(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);
+		iommu_domain_free(domain);
+		return -EINVAL;
+	}
+out:
+	etiommu->domains[0] = domain;
+	return 0;
+}
+
+/* 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 *edgetpu_pdev = to_abrolhos_dev(etdev);
+#endif
+	struct edgetpu_iommu *etiommu;
+	int ret;
+
+	etiommu = kzalloc(sizeof(*etiommu), GFP_KERNEL);
+	if (!etiommu)
+		return -ENOMEM;
+	idr_init(&etiommu->domain_pool);
+	mutex_init(&etiommu->pool_lock);
+	etiommu->iommu_group = iommu_group_get(etdev->dev);
+	if (etiommu->iommu_group) {
+		iommu_group_set_name(etiommu->iommu_group, "edgetpu");
+		dev_dbg(etdev->dev, "iommu group id %d setup\n",
+			iommu_group_id(etiommu->iommu_group));
+	} else {
+		dev_warn(etdev->dev, "device has no iommu group\n");
+	}
+
+	iommu_dev_enable_feature(etdev->dev, IOMMU_DEV_FEAT_AUX);
+	if (!iommu_dev_feature_enabled(etdev->dev, IOMMU_DEV_FEAT_AUX))
+		etdev_warn(etdev, "AUX domains not supported\n");
+	else
+		etiommu->aux_enabled = true;
+	ret = check_default_domain(etdev, etiommu);
+	if (ret)
+		goto err_free;
+
+	ret = edgetpu_register_iommu_device_fault_handler(etdev);
+	if (ret)
+		etdev_warn(etdev, "Failed to register fault handler! (%d)\n",
+			   ret);
+
+	/* etiommu initialization done */
+	etdev->mmu_cookie = etiommu;
+	/* TODO (b/178571278): remove chipset specific code. */
+#ifdef CONFIG_ABROLHOS
+	if (!edgetpu_pdev->csr_iova)
+		goto success;
+
+	etdev_dbg(etdev, "Mapping device CSRs: %llX -> %llX (%lu bytes)\n",
+		  edgetpu_pdev->csr_iova, edgetpu_pdev->csr_paddr,
+		  edgetpu_pdev->csr_size);
+
+	/* Add an IOMMU translation for the CSR region */
+	ret = edgetpu_mmu_add_translation(etdev, edgetpu_pdev->csr_iova,
+					  edgetpu_pdev->csr_paddr,
+					  edgetpu_pdev->csr_size,
+					  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:
+	kfree(etiommu);
+	return ret;
+}
+
+void edgetpu_mmu_reset(struct edgetpu_dev *etdev)
+{
+	/* If need to reset IOMMU driver can issue here. */
+}
+
+void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
+{
+#ifdef CONFIG_ABROLHOS
+	struct abrolhos_platform_dev *edgetpu_pdev = to_abrolhos_dev(etdev);
+#endif
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+	int i, ret;
+
+	if (!etiommu)
+		return;
+
+#ifdef CONFIG_ABROLHOS
+	if (edgetpu_pdev->csr_iova) {
+		edgetpu_mmu_remove_translation(&edgetpu_pdev->edgetpu_dev,
+					       edgetpu_pdev->csr_iova,
+					       edgetpu_pdev->csr_size,
+					       EDGETPU_CONTEXT_KCI);
+	}
+	edgetpu_pdev->csr_iova = 0;
+#endif
+	ret = edgetpu_unregister_iommu_device_fault_handler(etdev);
+	if (ret)
+		etdev_warn(etdev,
+			   "Failed to unregister device fault handler (%d)\n",
+			   ret);
+	edgetpu_mmu_reset(etdev);
+
+	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->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])
+		iommu_domain_free(etiommu->domains[0]);
+
+	idr_for_each(&etiommu->domain_pool, edgetpu_idr_free_domain_callback,
+		     NULL);
+	idr_destroy(&etiommu->domain_pool);
+	kfree(etiommu);
+	etdev->mmu_cookie = NULL;
+}
+
+int edgetpu_mmu_reattach(struct edgetpu_dev *etdev)
+{
+	return 0;
+}
+
+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)
+{
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+	size_t size = 0;
+	int prot = __dma_dir_to_iommu_prot(map->dir);
+	struct iommu_domain *domain;
+	int i;
+	struct scatterlist *sg;
+
+	if (!etiommu)
+		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;
+	return 0;
+}
+
+int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
+		    enum edgetpu_context_id context_id, u32 mmu_flags)
+{
+	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);
+
+	if (ret)
+		return ret;
+
+	if (mmu_flags & EDGETPU_MMU_64)
+		dev_warn_once(etdev->dev,
+			      "%s: 64-bit addressing is not supported",
+			      __func__);
+
+	ret = dma_map_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.nents,
+			       edgetpu_host_dma_dir(map->dir), map->dma_attrs);
+	if (!ret)
+		return -EINVAL;
+	map->sgt.nents = ret;
+	iova = sg_dma_address(map->sgt.sgl);
+
+	/*
+	 * 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) {
+		if (!iommu_map_sg(params.domain, iova, map->sgt.sgl,
+				  map->sgt.orig_nents, params.prot)) {
+			/* Undo the mapping in the default domain */
+			dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl,
+					   map->sgt.orig_nents,
+					   edgetpu_host_dma_dir(map->dir),
+					   DMA_ATTR_SKIP_CPU_SYNC);
+			return -ENOMEM;
+		}
+	}
+
+	map->device_address = iova;
+	return 0;
+}
+
+void edgetpu_mmu_unmap(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
+		       enum edgetpu_context_id context_id)
+{
+	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);
+	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);
+	}
+
+	/* Undo the mapping in the default domain */
+	dma_unmap_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.orig_nents,
+			   edgetpu_host_dma_dir(map->dir), map->dma_attrs);
+}
+
+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 tpu_addr_t orig_iova = iova;
+	struct scatterlist *sg;
+	int i;
+	int ret;
+
+	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i) {
+		ret = edgetpu_mmu_add_translation(etdev, iova, sg_phys(sg),
+						  sg->length, prot, context_id);
+		if (ret)
+			goto error;
+		iova += sg->length;
+	}
+	return 0;
+
+error:
+	edgetpu_mmu_remove_translation(etdev, orig_iova, iova - orig_iova,
+				       context_id);
+	return ret;
+}
+
+void edgetpu_mmu_unmap_iova_sgt_attrs(struct edgetpu_dev *etdev,
+				      tpu_addr_t iova, struct sg_table *sgt,
+				      enum dma_data_direction dir,
+				      enum edgetpu_context_id context_id,
+				      unsigned long attrs)
+{
+	size_t size = 0;
+	struct scatterlist *sg;
+	int i;
+
+	for_each_sg(sgt->sgl, sg, sgt->orig_nents, i)
+		size += sg->length;
+	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;
+
+	domain = get_domain_by_context_id(etdev, context_id);
+	if (!domain)
+		return -ENODEV;
+	return iommu_map(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;
+
+	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 = __dma_dir_to_iommu_prot(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 = __dma_dir_to_iommu_prot(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));
+}
+
+void edgetpu_mmu_use_dev_dram(struct edgetpu_dev *etdev, bool use_dev_dram)
+{
+}
+
+/* to be returned when domain aux is not supported */
+static struct edgetpu_iommu_domain invalid_etdomain = {
+	.pasid = IOMMU_PASID_INVALID,
+	.token = EDGETPU_DOMAIN_TOKEN_END,
+};
+
+struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_iommu_domain *etdomain =
+		kzalloc(sizeof(*etdomain), GFP_KERNEL);
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+	struct iommu_domain *domain;
+	int token;
+
+	if (!etdomain)
+		return NULL;
+	if (!etiommu->aux_enabled)
+		return &invalid_etdomain;
+	domain = iommu_domain_alloc(etdev->dev->bus);
+	if (!domain) {
+		etdev_warn(etdev, "iommu domain alloc failed");
+		return NULL;
+	}
+
+	mutex_lock(&etiommu->pool_lock);
+	token = idr_alloc(&etiommu->domain_pool, domain, 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);
+		iommu_domain_free(domain);
+		return NULL;
+	}
+	edgetpu_init_etdomain(etdomain, domain, token);
+	return etdomain;
+}
+
+void edgetpu_mmu_free_domain(struct edgetpu_dev *etdev,
+			     struct edgetpu_iommu_domain *etdomain)
+{
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+
+	if (!etdomain || etdomain == &invalid_etdomain)
+		return;
+	if (etdomain->pasid != IOMMU_PASID_INVALID) {
+		etdev_warn(etdev, "Domain should be detached before free");
+		edgetpu_mmu_detach_domain(etdev, etdomain);
+	}
+	mutex_lock(&etiommu->pool_lock);
+	idr_remove(&etiommu->domain_pool, etdomain->token);
+	mutex_unlock(&etiommu->pool_lock);
+	iommu_domain_free(etdomain->iommu_domain);
+	kfree(etdomain);
+}
+
+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;
+	int ret;
+	uint pasid;
+
+	/* Changes nothing if domain AUX is not supported. */
+	if (!etiommu->aux_enabled)
+		return 0;
+	if (etdomain->pasid != IOMMU_PASID_INVALID)
+		return -EINVAL;
+	domain = etdomain->iommu_domain;
+	ret = iommu_aux_attach_device(domain, etdev->dev);
+	if (ret) {
+		etdev_warn(etdev, "Attach IOMMU aux failed: %d", ret);
+		return ret;
+	}
+	pasid = iommu_aux_get_pasid(domain, etdev->dev);
+	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS) {
+		etdev_warn(etdev, "Invalid PASID %d returned from iommu",
+			   pasid);
+		ret = -EINVAL;
+		goto err_detach;
+	}
+	/* the IOMMU driver returned a duplicate PASID */
+	if (etiommu->domains[pasid]) {
+		ret = -EBUSY;
+		goto err_detach;
+	}
+	etiommu->domains[pasid] = domain;
+	etdomain->pasid = pasid;
+	return 0;
+err_detach:
+	iommu_aux_detach_device(domain, etdev->dev);
+	return ret;
+}
+
+void edgetpu_mmu_detach_domain(struct edgetpu_dev *etdev,
+			       struct edgetpu_iommu_domain *etdomain)
+{
+	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
+	uint pasid = etdomain->pasid;
+
+	if (!etiommu->aux_enabled)
+		return;
+	if (pasid <= 0 || pasid >= EDGETPU_NCONTEXTS)
+		return;
+	etiommu->domains[pasid] = NULL;
+	etdomain->pasid = IOMMU_PASID_INVALID;
+	iommu_aux_detach_device(etdomain->iommu_domain, etdev->dev);
+}
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 609d411..40f0b70 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -568,15 +568,6 @@ int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
 	return edgetpu_kci_send_cmd(kci, &cmd);
 }
 
-int edgetpu_kci_ack(struct edgetpu_kci *kci)
-{
-	struct edgetpu_command_element cmd = {
-		.code = KCI_CODE_ACK,
-	};
-
-	return edgetpu_kci_send_cmd(kci, &cmd);
-}
-
 int edgetpu_kci_map_log_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
 			       u32 size)
 {
@@ -688,8 +679,8 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 
 	/* If allocation failed still try handshake without full fw_info */
 	if (ret) {
-		etdev_warn(etdev, "%s: failed to allocate fw info buffer",
-			   __func__);
+		etdev_warn(etdev, "%s: error setting up fw info buffer: %d",
+			   __func__, ret);
 		memset(fw_info, 0, sizeof(*fw_info));
 	} else {
 		memset(mem.vaddr, 0, sizeof(*fw_info));
@@ -731,7 +722,7 @@ enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 	return flavor;
 }
 
-void edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
+int edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 {
 #define EDGETPU_USAGE_BUFFER_SIZE	4096
 	struct edgetpu_command_element cmd = {
@@ -751,22 +742,23 @@ void edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 	if (ret) {
 		etdev_warn_once(etdev, "%s: failed to allocate usage buffer",
 				__func__);
-		return;
+		return ret;
 	}
 
 	cmd.dma.address = mem.tpu_addr;
 	cmd.dma.size = EDGETPU_USAGE_BUFFER_SIZE;
-	memset(mem.vaddr, 0, sizeof(struct usage_tracker_header));
+	memset(mem.vaddr, 0, sizeof(struct edgetpu_usage_header));
 	ret = edgetpu_kci_send_cmd_return_resp(etdev->kci, &cmd, &resp);
 
 	if (ret == KCI_ERROR_UNIMPLEMENTED || ret == KCI_ERROR_UNAVAILABLE)
 		etdev_dbg(etdev, "firmware does not report usage\n");
 	else if (ret == KCI_ERROR_OK)
 		edgetpu_usage_stats_process_buffer(etdev, mem.vaddr);
-	else
+	else if (ret != -ETIMEDOUT)
 		etdev_warn_once(etdev, "%s: error %d", __func__, ret);
 
 	edgetpu_iremap_free(etdev, &mem, EDGETPU_CONTEXT_KCI);
+	return ret;
 }
 
 /* debugfs mappings dump */
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index aa77d9a..09189c0 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -224,14 +224,6 @@ int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
 			     u32 size, enum dma_data_direction dir);
 
 /*
- * Sends an ACK command and expects a response.
- * Can be used to test the firmware is running.
- *
- * Returns 0 if successful
- */
-int edgetpu_kci_ack(struct edgetpu_kci *kci);
-
-/*
  * Sends a FIRMWARE_INFO command and expects a response with a
  * edgetpu_fw_info struct filled out, including what firmware type is running,
  * along with build CL and time.
@@ -245,8 +237,13 @@ int edgetpu_kci_ack(struct edgetpu_kci *kci);
 enum edgetpu_fw_flavor edgetpu_kci_fw_info(
 	struct edgetpu_kci *kci, struct edgetpu_fw_info *fw_info);
 
-/* Retrieve usage tracking data from firmware, update info on host. */
-void edgetpu_kci_update_usage(struct edgetpu_dev *etdev);
+/*
+ * Retrieve usage tracking data from firmware, update info on host.
+ * Also used as a watchdog ping to firmware.
+ *
+ * Returns KCI response code on success or < 0 on error (typically -ETIMEDOUT).
+ */
+int edgetpu_kci_update_usage(struct edgetpu_dev *etdev);
 
 /*
  * Sends the "Map Log Buffer" command and waits for remote response.
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index b700c19..7d89754 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -188,7 +188,8 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 		}
 	}
 	/* Phase 2: Request state */
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL, state | PREQ);
+	edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL,
+			     (state << PSTATE_SHIFT) | PREQ);
 	SIM_PCHANNEL(etdev);
 
 	/* don't wait for state accept if STATE RUN */
@@ -200,7 +201,7 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 				     (val & PACCEPT) || (val & PDENY));
 	if (val & PDENY) {
 		edgetpu_dev_write_32(etdev, EDGETPU_REG_POWER_CONTROL,
-				     val & !state);
+				     val & ~(state << PSTATE_SHIFT));
 		etdev_dbg(etdev, "p-channel state change request denied\n");
 		deny = true;
 	}
diff --git a/drivers/edgetpu/edgetpu-sw-watchdog.c b/drivers/edgetpu/edgetpu-sw-watchdog.c
index 2db2b47..703a5e3 100644
--- a/drivers/edgetpu/edgetpu-sw-watchdog.c
+++ b/drivers/edgetpu/edgetpu-sw-watchdog.c
@@ -55,9 +55,9 @@ static void sw_wdt_work(struct work_struct *work)
 		container_of(dwork, struct edgetpu_sw_wdt, dwork);
 	struct edgetpu_dev *etdev = etdev_sw_wdt->etdev;
 
-	/* ping f/w */
+	/* Ping f/w, and grab updated usage stats while we're at it. */
 	etdev_dbg(etdev, "sw wdt: pinging firmware\n");
-	ret = edgetpu_kci_ack(etdev->kci);
+	ret = edgetpu_kci_update_usage(etdev);
 	if (ret)
 		etdev_dbg(etdev, "sw-watchdog ping resp:%d\n", ret);
 	if (ret == -ETIMEDOUT) {
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index 3b072d6..576274b 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -9,6 +9,7 @@
 #include <linux/sysfs.h>
 
 #include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
 #include "edgetpu-usage-stats.h"
 
 #if IS_ENABLED(CONFIG_ABROLHOS)
@@ -22,11 +23,18 @@ static enum tpu_pwr_state tpu_states_arr[] = {
 	TPU_ACTIVE_OD,
 };
 
-#else /* !CONFIG_ABROLHOS */
+#else /* CONFIG_HERMOSA */
 
-/* All execution times will be added to the same state. */
 static uint32_t tpu_states_arr[] = {
-	0,
+	4,	/* kActiveMinPower, kActiveVeryLowPower: 400MHz */
+	5,	/* kActiveLowPower: 800MHz */
+	6,	/* kActive: 950MHz */
+};
+
+static uint32_t tpu_states_display[] = {
+	400,
+	800,
+	950,
 };
 
 #endif /* CONFIG_ABROLHOS */
@@ -105,26 +113,50 @@ int edgetpu_usage_add(struct edgetpu_dev *etdev, struct tpu_usage *tpu_usage)
 	return 0;
 }
 
+static void edgetpu_utilization_update(
+	struct edgetpu_dev *etdev,
+	struct edgetpu_component_activity *activity)
+{
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+
+	if (!ustats)
+		return;
+
+	etdev_dbg(etdev, "%s: comp=%d utilized %d%%\n", __func__,
+		  activity->component, activity->utilization);
+
+	mutex_lock(&ustats->usage_stats_lock);
+	if (activity->utilization && activity->component >= 0 &&
+	    activity->component < EDGETPU_USAGE_COMPONENT_COUNT)
+		ustats->component_utilization[activity->component] =
+			activity->utilization;
+	mutex_unlock(&ustats->usage_stats_lock);
+}
+
 void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 {
-	struct usage_tracker_header *header = buf;
-	struct usage_tracker_metric *metric =
-		(struct usage_tracker_metric *)(header + 1);
+	struct edgetpu_usage_header *header = buf;
+	struct edgetpu_usage_metric *metric =
+		(struct edgetpu_usage_metric *)(header + 1);
 	int i;
 
 	etdev_dbg(etdev, "%s: n=%u sz=%u", __func__,
 		  header->num_metrics, header->metric_size);
-	if (header->metric_size != sizeof(struct usage_tracker_metric)) {
+	if (header->metric_size != sizeof(struct edgetpu_usage_metric)) {
 		etdev_dbg(etdev, "%s: expected sz=%zu, discard", __func__,
-			  sizeof(struct usage_tracker_metric));
+			  sizeof(struct edgetpu_usage_metric));
 		return;
 	}
 
 	for (i = 0; i < header->num_metrics; i++) {
 		switch (metric->type) {
-		case metric_type_tpu_usage:
+		case EDGETPU_METRIC_TYPE_TPU_USAGE:
 			edgetpu_usage_add(etdev, &metric->tpu_usage);
 			break;
+		case EDGETPU_METRIC_TYPE_COMPONENT_ACTIVITY:
+			edgetpu_utilization_update(
+				etdev, &metric->component_activity);
+			break;
 		default:
 			etdev_dbg(etdev, "%s: %d: skip unknown type=%u",
 				  __func__, i, metric->type);
@@ -146,12 +178,17 @@ static ssize_t tpu_usage_show(struct device *dev,
 	unsigned int bkt;
 	struct uid_entry *uid_entry;
 
-	/* uid: TPU_ACTIVE_SUD TPU_ACTIVE_UD TPU_ACTIVE_NOM TPU_ACTIVE_OD */
+	edgetpu_kci_update_usage(etdev);
+	/* uid: state0speed state1speed ... */
 	ret += scnprintf(buf, PAGE_SIZE, "uid:");
 
 	for (i = 0; i < NUM_TPU_STATES; i++)
 		ret += scnprintf(buf + ret, PAGE_SIZE - ret, " %d",
+#if IS_ENABLED(CONFIG_HERMOSA)
+				 tpu_states_display[i]);
+#else
 				 tpu_states_arr[i]);
+#endif
 
 	ret += scnprintf(buf + ret, PAGE_SIZE - ret, "\n");
 
@@ -205,6 +242,50 @@ static ssize_t tpu_usage_clear(struct device *dev,
 
 static DEVICE_ATTR(tpu_usage, 0644, tpu_usage_show, tpu_usage_clear);
 
+static ssize_t device_utilization_show(struct device *dev,
+				       struct device_attribute *attr,
+				       char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int32_t val;
+
+	edgetpu_kci_update_usage(etdev);
+	mutex_lock(&ustats->usage_stats_lock);
+	val = ustats->component_utilization[EDGETPU_USAGE_COMPONENT_DEVICE];
+	ustats->component_utilization[EDGETPU_USAGE_COMPONENT_DEVICE] = 0;
+	mutex_unlock(&ustats->usage_stats_lock);
+	return scnprintf(buf, PAGE_SIZE, "%d\n", val);
+}
+static DEVICE_ATTR_RO(device_utilization);
+
+static ssize_t tpu_utilization_show(struct device *dev,
+				    struct device_attribute *attr,
+				    char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	int32_t val;
+
+	edgetpu_kci_update_usage(etdev);
+	mutex_lock(&ustats->usage_stats_lock);
+	val = ustats->component_utilization[EDGETPU_USAGE_COMPONENT_TPU];
+	ustats->component_utilization[EDGETPU_USAGE_COMPONENT_TPU] = 0;
+	mutex_unlock(&ustats->usage_stats_lock);
+	return scnprintf(buf, PAGE_SIZE, "%d\n", val);
+}
+static DEVICE_ATTR_RO(tpu_utilization);
+
+static struct attribute *usage_stats_dev_attrs[] = {
+	&dev_attr_tpu_usage.attr,
+	&dev_attr_device_utilization.attr,
+	&dev_attr_tpu_utilization.attr,
+	NULL,
+};
+
+static const struct attribute_group usage_stats_attr_group = {
+	.attrs = usage_stats_dev_attrs,
+};
 void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_usage_stats *ustats;
@@ -220,12 +301,11 @@ void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 
 	hash_init(ustats->uid_hash_table);
 	mutex_init(&ustats->usage_stats_lock);
-
 	etdev->usage_stats = ustats;
 
-	ret = device_create_file(etdev->dev, &dev_attr_tpu_usage);
+	ret = device_add_group(etdev->dev, &usage_stats_attr_group);
 	if (ret)
-		etdev_warn(etdev, "failed to create the usage_stats file\n");
+		etdev_warn(etdev, "failed to create the usage_stats attrs\n");
 
 	etdev_dbg(etdev, "%s init\n", __func__);
 }
@@ -236,7 +316,7 @@ void edgetpu_usage_stats_exit(struct edgetpu_dev *etdev)
 
 	if (ustats) {
 		usage_stats_remove_uids(ustats);
-		device_remove_file(etdev->dev, &dev_attr_tpu_usage);
+		device_remove_group(etdev->dev, &usage_stats_attr_group);
 	}
 
 	etdev_dbg(etdev, "%s exit\n", __func__);
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 68c6539..6b72747 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -12,7 +12,7 @@
 
 /* Header struct in the metric buffer. */
 /* Must be kept in sync with firmware struct UsageTrackerHeader */
-struct usage_tracker_header {
+struct edgetpu_usage_header {
 	uint32_t num_metrics;		/* Number of metrics being reported */
 	uint32_t metric_size;		/* Size of each metric struct */
 };
@@ -31,21 +31,45 @@ struct tpu_usage {
 	uint32_t duration_us;
 };
 
+/*
+ * An enum to represent the different components we can track metrics for.
+ * Must be kept in sync with firmware struct Component.
+ */
+enum edgetpu_usage_component {
+	/* The device as a whole (TPU, R52, DMA330, etc.) */
+	EDGETPU_USAGE_COMPONENT_DEVICE = 0,
+	/* Just the TPU core (scalar core and tiles) */
+	EDGETPU_USAGE_COMPONENT_TPU = 1,
+	EDGETPU_USAGE_COMPONENT_COUNT = 2, /* number of components above */
+};
+
+/*
+ * Encapsulates information about activity of a component.
+ * Must be kept in sync with firmware struct ComponentActivity.
+ */
+struct edgetpu_component_activity {
+	enum edgetpu_usage_component component;
+	/* Utilization as a percentage since the last read. */
+	int32_t utilization;
+};
+
 /* Must be kept in sync with firmware enum class UsageTrackerMetric::Type */
-enum usage_tracker_metric_type {
-	metric_type_reserved = 0,
-	metric_type_tpu_usage = 1,
+enum edgetpu_usage_metric_type {
+	EDGETPU_METRIC_TYPE_RESERVED = 0,
+	EDGETPU_METRIC_TYPE_TPU_USAGE = 1,
+	EDGETPU_METRIC_TYPE_COMPONENT_ACTIVITY = 2,
 };
 
 /*
  * Encapsulates a single metric reported to the kernel.
  * Must be kept in sync with firmware struct UsageTrackerMetric.
  */
-struct usage_tracker_metric {
+struct edgetpu_usage_metric {
 	uint32_t type;
 	uint8_t reserved[4];
 	union {
 		struct tpu_usage tpu_usage;
+		struct edgetpu_component_activity component_activity;
 	};
 };
 
@@ -53,6 +77,8 @@ struct usage_tracker_metric {
 
 struct edgetpu_usage_stats {
 	DECLARE_HASHTABLE(uid_hash_table, UID_HASH_BITS);
+	/* component utilization values reported by firmware */
+	int32_t component_utilization[EDGETPU_USAGE_COMPONENT_COUNT];
 	struct mutex usage_stats_lock;
 };