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

* whitechapel:
  edgetpu: use better defines to select right code
  edgetpu: tpu_usage_add take struct tpu_usage as parameter
  edgetpu: set ctx to domain token when mbox detached
  edgetpu: abrolhos get domain by context ID
  edgetpu: abrolhos support iommu token
  edgetpu: introduce token to edgetpu iommu domain
  edgetpu: abrolhos fix mmu attach
  edgetpu: abrolhos fix mmu_attach potential mem leak

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

10 files changed, 171 insertions, 91 deletions

diff --git a/drivers/edgetpu/abrolhos-iommu.c b/drivers/edgetpu/abrolhos-iommu.c
index d776a63..097d576 100644
--- a/drivers/edgetpu/abrolhos-iommu.c
+++ b/drivers/edgetpu/abrolhos-iommu.c
@@ -7,6 +7,7 @@
 
 #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>
@@ -25,6 +26,12 @@ struct edgetpu_iommu {
 	 * 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;
 };
@@ -35,13 +42,42 @@ struct edgetpu_iommu_map_params {
 	struct iommu_domain *domain;
 };
 
-static struct iommu_domain *get_domain_by_pasid(struct edgetpu_dev *etdev,
-						uint pasid)
+/*
+ * 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];
 
@@ -114,6 +150,15 @@ edgetpu_unregister_iommu_device_fault_handler(struct edgetpu_dev *etdev)
 
 #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)
@@ -129,10 +174,11 @@ static int edgetpu_iommu_fault_handler(struct iommu_domain *domain,
 
 static void edgetpu_init_etdomain(struct edgetpu_iommu_domain *etdomain,
 				  struct iommu_domain *domain,
-				  unsigned int pasid)
+				  int token)
 {
 	etdomain->iommu_domain = domain;
-	etdomain->pasid = pasid;
+	etdomain->pasid = IOMMU_PASID_INVALID;
+	etdomain->token = token;
 	iommu_set_fault_handler(domain, edgetpu_iommu_fault_handler, etdomain);
 }
 
@@ -193,8 +239,8 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	etiommu = kzalloc(sizeof(*etiommu), GFP_KERNEL);
 	if (!etiommu)
 		return -ENOMEM;
-
-	etdev->mmu_cookie = etiommu;
+	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");
@@ -211,15 +257,17 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 		etiommu->aux_enabled = true;
 	ret = check_default_domain(etdev, etiommu);
 	if (ret)
-		return 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)
-		return 0;
+		goto success;
 
 	etdev_dbg(etdev, "Mapping device CSRs: %llX -> %llX (%lu bytes)\n",
 		  edgetpu_pdev->csr_iova, edgetpu_pdev->csr_paddr,
@@ -234,10 +282,16 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	if (ret) {
 		etdev_err(etdev, "Unable to map device CSRs into IOMMU\n");
 		edgetpu_unregister_iommu_device_fault_handler(etdev);
-		return ret;
+		etdev->mmu_cookie = NULL;
+		goto err_free;
 	}
 
+success:
 	return 0;
+
+err_free:
+	kfree(etiommu);
+	return ret;
 }
 
 void edgetpu_mmu_reset(struct edgetpu_dev *etdev)
@@ -271,16 +325,21 @@ void edgetpu_mmu_detach(struct edgetpu_dev *etdev)
 
 	for (i = etiommu->context_0_default ? 1 : 0; i < EDGETPU_NCONTEXTS;
 	     i++) {
-		if (etiommu->domains[i]) {
+		if (etiommu->domains[i])
 			iommu_aux_detach_device(etiommu->domains[i],
 						etdev->dev);
-			iommu_domain_free(etiommu->domains[i]);
-		}
 	}
 
 	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;
 }
@@ -290,12 +349,6 @@ int edgetpu_mmu_reattach(struct edgetpu_dev *etdev)
 	return 0;
 }
 
-/* Return context ID enumeration value as a Process Address Space ID. */
-static uint context_id_to_pasid(enum edgetpu_context_id context_id)
-{
-	return (uint)context_id;
-}
-
 static int get_iommu_map_params(struct edgetpu_dev *etdev,
 				struct edgetpu_mapping *map,
 				enum edgetpu_context_id context_id,
@@ -303,20 +356,15 @@ static int get_iommu_map_params(struct edgetpu_dev *etdev,
 {
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	size_t size = 0;
-	uint pasid = context_id_to_pasid(context_id);
 	int prot = __dma_dir_to_iommu_prot(map->dir);
 	struct iommu_domain *domain;
 	int i;
 	struct scatterlist *sg;
 
-	if (pasid >= EDGETPU_NCONTEXTS) {
-		etdev_err(etdev, "Invalid context_id %d\n", context_id);
-		return -EINVAL;
-	}
 	if (!etiommu)
 		return -EINVAL;
 
-	domain = get_domain_by_pasid(etdev, pasid);
+	domain = get_domain_by_context_id(etdev, context_id);
 	if (!domain) {
 		etdev_err(etdev, "Unable to find an iommu domain\n");
 		return -ENODEV;
@@ -396,10 +444,6 @@ void edgetpu_mmu_unmap(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
 		iommu_unmap(params.domain, map->device_address, params.size);
 	}
 
-	/*
-	 * Always do dma_unmap since context_id might be invalid when group has
-	 * mailbox detached.
-	 */
 	/* 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);
@@ -466,12 +510,8 @@ int edgetpu_mmu_add_translation(struct edgetpu_dev *etdev, unsigned long iova,
 				enum edgetpu_context_id context_id)
 {
 	struct iommu_domain *domain;
-	uint pasid = context_id_to_pasid(context_id);
 
-	if (pasid >= EDGETPU_NCONTEXTS)
-		return -EINVAL;
-
-	domain = get_domain_by_pasid(etdev, pasid);
+	domain = get_domain_by_context_id(etdev, context_id);
 	if (!domain)
 		return -ENODEV;
 	return iommu_map(domain, iova, paddr, size, prot);
@@ -482,9 +522,8 @@ void edgetpu_mmu_remove_translation(struct edgetpu_dev *etdev,
 				    enum edgetpu_context_id context_id)
 {
 	struct iommu_domain *domain;
-	uint pasid = context_id_to_pasid(context_id);
 
-	domain = get_domain_by_pasid(etdev, pasid);
+	domain = get_domain_by_context_id(etdev, context_id);
 	if (domain)
 		iommu_unmap(domain, iova, size);
 }
@@ -499,12 +538,8 @@ tpu_addr_t edgetpu_mmu_tpu_map(struct edgetpu_dev *etdev, dma_addr_t down_addr,
 		iommu_get_domain_for_dev(etdev->dev);
 	phys_addr_t paddr;
 	int prot = __dma_dir_to_iommu_prot(dir);
-	uint pasid = context_id_to_pasid(context_id);
-
-	if (pasid >= EDGETPU_NCONTEXTS)
-		return 0;
-	domain = get_domain_by_pasid(etdev, pasid);
 
+	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
@@ -529,10 +564,8 @@ void edgetpu_mmu_tpu_unmap(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
 	struct iommu_domain *domain;
 	struct iommu_domain *default_domain =
 		iommu_get_domain_for_dev(etdev->dev);
-	uint pasid = context_id_to_pasid(context_id);
-
-	domain = get_domain_by_pasid(etdev, pasid);
 
+	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
@@ -551,6 +584,7 @@ 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)
@@ -559,6 +593,7 @@ struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
 		kzalloc(sizeof(*etdomain), GFP_KERNEL);
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	struct iommu_domain *domain;
+	int token;
 
 	if (!etdomain)
 		return NULL;
@@ -570,19 +605,33 @@ struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
 		return NULL;
 	}
 
-	edgetpu_init_etdomain(etdomain, domain, IOMMU_PASID_INVALID);
+	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);
 }
diff --git a/drivers/edgetpu/abrolhos/config.h b/drivers/edgetpu/abrolhos/config.h
index f91ddbc..262b3a5 100644
--- a/drivers/edgetpu/abrolhos/config.h
+++ b/drivers/edgetpu/abrolhos/config.h
@@ -12,6 +12,8 @@
 
 #define EDGETPU_DEV_MAX		1
 
+#define EDGETPU_HAS_MULTI_GROUPS
+
 /*
  * A remapped data region is available. This will be accessible by the R52
  * regardless of active context and is typically used for logging buffer and
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index c5f2089..5a13adb 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -14,6 +14,12 @@
 
 #else /* !CONFIG_HERMOSA */
 
+#ifdef CONFIG_JANEIRO
+
+#include "janeiro/config.h"
+
+#else
+
 #ifndef CONFIG_ABROLHOS
 #define CONFIG_ABROLHOS
 #warning "Building default chipset abrolhos"
@@ -21,6 +27,7 @@
 
 #include "abrolhos/config.h"
 
+#endif /* CONFIG_JANEIRO */
 #endif /* CONFIG_HERMOSA */
 
 #define EDGETPU_DEFAULT_FIRMWARE_NAME "google/edgetpu-" DRIVER_NAME ".fw"
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index f323a3c..d42da25 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -59,7 +59,7 @@ struct edgetpu_host_map {
 	struct sg_table *sg_tables;
 };
 
-#if !IS_ENABLED(CONFIG_ABROLHOS)
+#ifdef EDGETPU_HAS_MCP
 
 /* parameter to be used in async KCI jobs */
 struct kci_worker_param {
@@ -90,7 +90,7 @@ static int edgetpu_kci_leave_group_worker(struct kci_worker_param *param)
 	return 0;
 }
 
-#endif /* CONFIG_ABROLHOS */
+#endif /* EDGETPU_HAS_MCP */
 
 static int edgetpu_group_kci_open_device(struct edgetpu_device_group *group)
 {
@@ -141,10 +141,10 @@ static void edgetpu_group_kci_close_device(struct edgetpu_device_group *group)
  */
 static void edgetpu_device_group_kci_leave(struct edgetpu_device_group *group)
 {
-#if IS_ENABLED(CONFIG_ABROLHOS)
+#ifdef EDGETPU_HAS_MULTI_GROUPS
 	edgetpu_kci_update_usage(group->etdev);
 	return edgetpu_group_kci_close_device(group);
-#else /* !CONFIG_ABROLHOS */
+#else /* !EDGETPU_HAS_MULTI_GROUPS */
 	struct kci_worker_param *params =
 		kmalloc_array(group->n_clients, sizeof(*params), GFP_KERNEL);
 	struct edgetpu_async_ctx *ctx = edgetpu_async_alloc_ctx();
@@ -173,7 +173,7 @@ static void edgetpu_device_group_kci_leave(struct edgetpu_device_group *group)
 out_free:
 	edgetpu_async_free_ctx(ctx);
 	kfree(params);
-#endif /* CONFIG_ABROLHOS */
+#endif /* EDGETPU_HAS_MULTI_GROUPS */
 }
 
 /*
@@ -184,9 +184,9 @@ out_free:
 static int
 edgetpu_device_group_kci_finalized(struct edgetpu_device_group *group)
 {
-#if IS_ENABLED(CONFIG_ABROLHOS)
+#ifdef EDGETPU_HAS_MULTI_GROUPS
 	return edgetpu_group_kci_open_device(group);
-#else /* !CONFIG_ABROLHOS */
+#else /* !EDGETPU_HAS_MULTI_GROUPS */
 	struct kci_worker_param *params =
 		kmalloc_array(group->n_clients, sizeof(*params), GFP_KERNEL);
 	struct edgetpu_async_ctx *ctx = edgetpu_async_alloc_ctx();
@@ -255,7 +255,7 @@ out_free:
 	edgetpu_async_free_ctx(ctx);
 	kfree(params);
 	return ret;
-#endif /* CONFIG_ABROLHOS */
+#endif /* EDGETPU_HAS_MULTI_GROUPS */
 }
 
 /*
@@ -1190,7 +1190,7 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 
 	mutex_lock(&group->lock);
 	context_id = edgetpu_group_context_id_locked(group);
-	if (!edgetpu_group_finalized_and_attached(group)) {
+	if (!edgetpu_device_group_is_finalized(group)) {
 		ret = -EINVAL;
 		goto error_unlock_group;
 	}
@@ -1262,7 +1262,7 @@ int edgetpu_device_group_unmap(struct edgetpu_device_group *group,
 	int ret = 0;
 
 	mutex_lock(&group->lock);
-	if (!edgetpu_group_finalized_and_attached(group)) {
+	if (!edgetpu_device_group_is_finalized(group)) {
 		ret = -EINVAL;
 		goto unlock_group;
 	}
@@ -1305,7 +1305,7 @@ int edgetpu_device_group_sync_buffer(struct edgetpu_device_group *group,
 		return -EINVAL;
 
 	mutex_lock(&group->lock);
-	if (!edgetpu_group_finalized_and_attached(group)) {
+	if (!edgetpu_device_group_is_finalized(group)) {
 		ret = -EINVAL;
 		goto unlock_group;
 	}
@@ -1476,7 +1476,11 @@ void edgetpu_group_detach_mailbox(struct edgetpu_device_group *group)
 	edgetpu_device_group_put(mailbox->internal.group);
 	edgetpu_mailbox_remove(mgr, mailbox);
 	edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
-	group->context_id = EDGETPU_CONTEXT_INVALID;
+	if (group->etdomain->token != EDGETPU_DOMAIN_TOKEN_END)
+		group->context_id =
+			EDGETPU_CONTEXT_DOMAIN_TOKEN | group->etdomain->token;
+	else
+		group->context_id = EDGETPU_CONTEXT_INVALID;
 
 	mutex_unlock(&group->lock);
 }
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 87a0987..cd47738 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -85,8 +85,9 @@ struct edgetpu_device_group {
 	/*
 	 * Context ID ranges from EDGETPU_CONTEXT_VII_BASE to
 	 * EDGETPU_NCONTEXTS - 1.
-	 * This equals EDGETPU_CONTEXT_INVALID when the group has mailbox
-	 * detached (means the group isn't in any context at this time).
+	 * This equals EDGETPU_CONTEXT_INVALID or a token OR'ed with
+	 * EDGETPU_CONTEXT_DOMAIN_TOKEN when the group has mailbox detached
+	 * (means the group isn't in any context at this time).
 	 */
 	enum edgetpu_context_id context_id;
 	/* The IOMMU domain being associated to this group */
@@ -351,7 +352,8 @@ int edgetpu_group_attach_mailbox(struct edgetpu_device_group *group);
 static inline bool
 edgetpu_group_mailbox_detached_locked(const struct edgetpu_device_group *group)
 {
-	return group->context_id == EDGETPU_CONTEXT_INVALID;
+	return group->context_id == EDGETPU_CONTEXT_INVALID ||
+	       group->context_id & EDGETPU_CONTEXT_DOMAIN_TOKEN;
 }
 
 /*
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index b424463..80bb853 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -749,7 +749,7 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 		goto err_put;
 
 	mutex_lock(&group->lock);
-	if (!edgetpu_group_finalized_and_attached(group))
+	if (!edgetpu_device_group_is_finalized(group))
 		goto err_unlock_group;
 
 	dmap = alloc_dmabuf_map(group, flags);
@@ -819,7 +819,7 @@ int edgetpu_unmap_dmabuf(struct edgetpu_device_group *group, u32 die_index,
 
 	mutex_lock(&group->lock);
 	/* the group is disbanded means all the mappings have been released */
-	if (!edgetpu_group_finalized_and_attached(group))
+	if (!edgetpu_device_group_is_finalized(group))
 		goto out_unlock;
 	edgetpu_mapping_lock(mappings);
 	map = edgetpu_mapping_find_locked(mappings, die_index, tpu_addr);
@@ -855,7 +855,7 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 	if (arg->size == 0)
 		return -EINVAL;
 	mutex_lock(&group->lock);
-	if (!edgetpu_group_finalized_and_attached(group))
+	if (!edgetpu_device_group_is_finalized(group))
 		goto err_unlock_group;
 	/* checks not all FDs are ignored */
 	for (i = 0; i < group->n_clients; i++)
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 237fdd2..8974125 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -74,6 +74,8 @@ enum edgetpu_context_id {
 	EDGETPU_CONTEXT_KCI = 0,	/* TPU firmware/kernel ID 0 */
 	EDGETPU_CONTEXT_VII_BASE = 1,	/* groups 0-6 IDs 1-7 */
 	/* contexts 8 and above not yet allocated */
+	/* A bit mask to mark the context is an IOMMU domain token */
+	EDGETPU_CONTEXT_DOMAIN_TOKEN = 1 << 30,
 };
 
 typedef u64 tpu_addr_t;
diff --git a/drivers/edgetpu/edgetpu-mmu.h b/drivers/edgetpu/edgetpu-mmu.h
index 8f453b2..0c53589 100644
--- a/drivers/edgetpu/edgetpu-mmu.h
+++ b/drivers/edgetpu/edgetpu-mmu.h
@@ -41,9 +41,24 @@
 #define EDGETPU_MMU_DEVICE	(1 << 2)
 #define EDGETPU_MMU_DMABUF	(2 << 2)
 
+/*
+ * The max possible value of token is (EDGETPU_DOMAIN_TOKEN_END - 1), which
+ * shouldn't equal or exceed the bit mask EDGETPU_CONTEXT_DOMAIN_TOKEN.
+ */
+#define EDGETPU_DOMAIN_TOKEN_END EDGETPU_CONTEXT_DOMAIN_TOKEN
 struct edgetpu_iommu_domain {
+	/*
+	 * IOMMU PASID, set by edgetpu_mmu_attach_domain().
+	 * This field should be set as IOMMU_PASID_INVALID in
+	 * edgetpu_mmu_detach_domain().
+	 */
 	uint pasid;
 	struct iommu_domain *iommu_domain;
+	/*
+	 * A token set by edgetpu_mmu_alloc_domain(). See the description of
+	 * edgetpu_mmu_add_translation() about @context_id for more details.
+	 */
+	int token;
 };
 
 /*
@@ -185,6 +200,12 @@ void edgetpu_mmu_free(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
  * may actually be another IOVA for another IOMMU downstream of the chip MMU
  * (as on Hermosa, where the SMMU translates TPU VAs to IOVAs sent to the IOMMU
  * downstream of the TPU).
+ *
+ * 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
+ * edgetpu_mmu_alloc_domain(). This description holds for all APIs in this file
+ * with @context_id as a parameter.
  */
 int edgetpu_mmu_add_translation(struct edgetpu_dev *etdev, unsigned long iova,
 				phys_addr_t paddr, size_t size, int prot,
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index b7b309c..3b072d6 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -65,23 +65,24 @@ find_uid_entry_locked(int32_t uid, struct edgetpu_usage_stats *ustats)
 	return NULL;
 }
 
-int edgetpu_usage_stats_add(int32_t uid, uint32_t state, uint32_t duration,
-			    struct edgetpu_dev *etdev)
+int edgetpu_usage_add(struct edgetpu_dev *etdev, struct tpu_usage *tpu_usage)
 {
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
 	struct uid_entry *uid_entry;
+	int state = tpu_state_map(tpu_usage->power_state);
 
 	if (!ustats)
 		return 0;
 
-	etdev_dbg(etdev, "%s: uid=%u state=%u dur=%u", __func__, uid, state,
-		  duration);
+	etdev_dbg(etdev, "%s: uid=%u state=%u dur=%u", __func__,
+		  tpu_usage->uid, tpu_usage->power_state,
+		  tpu_usage->duration_us);
 	mutex_lock(&ustats->usage_stats_lock);
 
 	/* Find the uid in uid_hash_table first */
-	uid_entry = find_uid_entry_locked(uid, ustats);
+	uid_entry = find_uid_entry_locked(tpu_usage->uid, ustats);
 	if (uid_entry) {
-		uid_entry->time_in_state[tpu_state_map(state)] += duration;
+		uid_entry->time_in_state[state] += tpu_usage->duration_us;
 		mutex_unlock(&ustats->usage_stats_lock);
 		return 0;
 	}
@@ -93,11 +94,11 @@ int edgetpu_usage_stats_add(int32_t uid, uint32_t state, uint32_t duration,
 		return -ENOMEM;
 	}
 
-	uid_entry->uid = uid;
-	uid_entry->time_in_state[tpu_state_map(state)] += duration;
+	uid_entry->uid = tpu_usage->uid;
+	uid_entry->time_in_state[state] += tpu_usage->duration_us;
 
 	/* Add uid_entry to the uid_hash_table */
-	hash_add(ustats->uid_hash_table, &uid_entry->node, uid);
+	hash_add(ustats->uid_hash_table, &uid_entry->node, tpu_usage->uid);
 
 	mutex_unlock(&ustats->usage_stats_lock);
 
@@ -122,14 +123,7 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 	for (i = 0; i < header->num_metrics; i++) {
 		switch (metric->type) {
 		case metric_type_tpu_usage:
-			{
-				struct tpu_usage *tpu_usage =
-					&metric->tpu_usage;
-
-				edgetpu_usage_stats_add(
-					tpu_usage->uid, tpu_usage->power_state,
-					tpu_usage->duration_us, etdev);
-			}
+			edgetpu_usage_add(etdev, &metric->tpu_usage);
 			break;
 		default:
 			etdev_dbg(etdev, "%s: %d: skip unknown type=%u",
@@ -141,9 +135,9 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 	}
 }
 
-static ssize_t usage_stats_show(struct device *dev,
-				struct device_attribute *attr,
-				char *buf)
+static ssize_t tpu_usage_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;
@@ -196,10 +190,10 @@ static void usage_stats_remove_uids(struct edgetpu_usage_stats *ustats)
 }
 
 /* Write to clear all entries in uid_hash_table */
-static ssize_t usage_stats_clear(struct device *dev,
-				 struct device_attribute *attr,
-				 const char *buf,
-				 size_t count)
+static ssize_t tpu_usage_clear(struct device *dev,
+			       struct device_attribute *attr,
+			       const char *buf,
+			       size_t count)
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
@@ -209,7 +203,7 @@ static ssize_t usage_stats_clear(struct device *dev,
 	return count;
 }
 
-static DEVICE_ATTR(usage_stats, 0644, usage_stats_show, usage_stats_clear);
+static DEVICE_ATTR(tpu_usage, 0644, tpu_usage_show, tpu_usage_clear);
 
 void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 {
@@ -229,7 +223,7 @@ void edgetpu_usage_stats_init(struct edgetpu_dev *etdev)
 
 	etdev->usage_stats = ustats;
 
-	ret = device_create_file(etdev->dev, &dev_attr_usage_stats);
+	ret = device_create_file(etdev->dev, &dev_attr_tpu_usage);
 	if (ret)
 		etdev_warn(etdev, "failed to create the usage_stats file\n");
 
@@ -242,7 +236,7 @@ void edgetpu_usage_stats_exit(struct edgetpu_dev *etdev)
 
 	if (ustats) {
 		usage_stats_remove_uids(ustats);
-		device_remove_file(etdev->dev, &dev_attr_usage_stats);
+		device_remove_file(etdev->dev, &dev_attr_tpu_usage);
 	}
 
 	etdev_dbg(etdev, "%s exit\n", __func__);
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 42d75df..68c6539 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -56,8 +56,7 @@ struct edgetpu_usage_stats {
 	struct mutex usage_stats_lock;
 };
 
-int edgetpu_usage_stats_add(int32_t uid, uint32_t state, uint32_t duration,
-			    struct edgetpu_dev *etdev);
+int edgetpu_usage_add(struct edgetpu_dev *etdev, struct tpu_usage *tpu_usage);
 void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf);
 void edgetpu_usage_stats_init(struct edgetpu_dev *etdev);
 void edgetpu_usage_stats_exit(struct edgetpu_dev *etdev);