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

Revert "gcip: temporary disable gcip-iommu"
gcip: temporary disable gcip-iommu
edgetpu: usage_stats send metrics v2 requests with v1 fallback
Bug: 271372136
edgetpu: pm: reject power up if thermal suspended
gcip: kci: add usage-stats metrics v1 / v2 commands
gcip: pm power_up callback add comments for thermal suspend suggestion
edgetpu: usage stats: sync additional metrics v2 changes
Bug: 271372136 (repeat)
edgetpu: usage stats add field definitions for metrics v2
gcip: implement a function returning default IOMMU domain
Bug: 243479562
gcip: implement map/unmap in legacy mode
Bug: 243479562 (repeat)
gcip: implement gcip_iommu_domain_{map,unmap}_sg
Bug: 243479562 (repeat)
gcip: implement gcip_iommu_domain_pool_{alloc,free}_domain
Bug: 243479562 (repeat)
gcip: add granule alignment functions
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_type
Bug: 243479562 (repeat)
gcip: implement funcs of gcip_iommu_domain_pool
Bug: 243479562 (repeat)
gcip: introduce a function returning default IOMMU domain
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_{map,unmap}_sg
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_pool_{alloc,free}_domain
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_ops
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_type
Bug: 243479562 (repeat)
gcip: introduce gcip_iommu_domain_pool and its funcs
Bug: 243479562 (repeat)

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

11 files changed, 815 insertions, 27 deletions

diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 7b9fbac..5e442d0 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -27,6 +27,7 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-sw-watchdog.h"
 #include "edgetpu-telemetry.h"
+#include "edgetpu-usage-stats.h"
 
 static char *firmware_name;
 module_param(firmware_name, charp, 0660);
@@ -339,6 +340,7 @@ static int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 				       enum edgetpu_firmware_flags flags)
 {
 	const struct edgetpu_firmware_chip_data *chip_fw = et_fw->p->chip_fw;
+	struct edgetpu_dev *etdev = et_fw->etdev;
 	struct edgetpu_firmware_desc new_fw_desc;
 	int ret;
 
@@ -349,7 +351,7 @@ static int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 	if (ret)
 		goto out_failed;
 
-	etdev_dbg(et_fw->etdev, "run fw %s flags=%#x", name, flags);
+	etdev_dbg(etdev, "run fw %s flags=%#x", name, flags);
 	if (chip_fw->prepare_run) {
 		ret = chip_fw->prepare_run(et_fw, &new_fw_desc.buf);
 		if (ret)
@@ -367,6 +369,9 @@ static int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 	if (!ret)
 		edgetpu_sw_wdt_start(et_fw->etdev);
 	edgetpu_firmware_set_state(et_fw, ret);
+	/* If previous firmware was metrics v1-only reset that flag and probe this again. */
+	if (etdev->usage_stats)
+		etdev->usage_stats->use_metrics_v1 = false;
 	return ret;
 
 out_unload_new_fw:
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index e57dddb..a3ecb11 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -519,7 +519,7 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 
 	if (ret) {
 		etdev_warn_ratelimited(client->etdev,
-				       "wakelock acquire rejected due to thermal suspend");
+		       "wakelock acquire rejected due to device thermal limit exceeded");
 		goto error_client_unlock;
 	}
 
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 9258aba..44d8542 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -480,7 +480,7 @@ int edgetpu_kci_update_usage_locked(struct edgetpu_dev *etdev)
 {
 #define EDGETPU_USAGE_BUFFER_SIZE	4096
 	struct gcip_kci_command_element cmd = {
-		.code = GCIP_KCI_CODE_GET_USAGE,
+		.code = GCIP_KCI_CODE_GET_USAGE_V2,
 		.dma = {
 			.address = 0,
 			.size = 0,
@@ -499,17 +499,26 @@ int edgetpu_kci_update_usage_locked(struct edgetpu_dev *etdev)
 		return ret;
 	}
 
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+retry_v1:
+	if (etdev->usage_stats && etdev->usage_stats->use_metrics_v1)
+		cmd.code = GCIP_KCI_CODE_GET_USAGE_V1;
 	cmd.dma.address = mem.tpu_addr;
 	cmd.dma.size = EDGETPU_USAGE_BUFFER_SIZE;
 	memset(mem.vaddr, 0, sizeof(struct edgetpu_usage_header));
 	ret = gcip_kci_send_cmd_return_resp(etdev->etkci->kci, &cmd, &resp);
 
-	if (ret == GCIP_KCI_ERROR_UNIMPLEMENTED || ret == GCIP_KCI_ERROR_UNAVAILABLE)
+	if (ret == GCIP_KCI_ERROR_UNIMPLEMENTED || ret == GCIP_KCI_ERROR_UNAVAILABLE) {
+		if (etdev->usage_stats && !etdev->usage_stats->use_metrics_v1) {
+			etdev->usage_stats->use_metrics_v1 = true;
+			goto retry_v1;
+		}
 		etdev_dbg(etdev, "firmware does not report usage\n");
-	else if (ret == GCIP_KCI_ERROR_OK)
+	} else if (ret == GCIP_KCI_ERROR_OK) {
 		edgetpu_usage_stats_process_buffer(etdev, mem.vaddr);
-	else if (ret != -ETIMEDOUT)
+	} else if (ret != -ETIMEDOUT) {
 		etdev_warn_once(etdev, "error %d", ret);
+	}
 
 	edgetpu_iremap_free(etdev, &mem, EDGETPU_CONTEXT_KCI);
 
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index db7a793..60751dd 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -241,23 +241,44 @@ out:
 
 void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 {
-	struct edgetpu_usage_header *header = buf;
-	struct edgetpu_usage_metric *metric =
-		(struct edgetpu_usage_metric *)(header + 1);
+	struct edgetpu_usage_stats *ustats = etdev->usage_stats;
+	struct edgetpu_usage_metric *metric;
+	uint metric_size;
+	uint num_metrics;
+	uint version;
 	int i;
 
-	etdev_dbg(etdev, "%s: n=%u sz=%u", __func__,
-		  header->num_metrics, header->metric_size);
-	if (header->metric_size < EDGETPU_USAGE_METRIC_SIZE_V1) {
+	if (!ustats)
+		return;
+
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+	if (ustats->use_metrics_v1) {
+		struct edgetpu_usage_header_v1 *header = buf;
+
+		metric_size = header->metric_size;
+		num_metrics = header->num_metrics;
+		version = 1;
+		metric = (struct edgetpu_usage_metric *)(header + 1);
+	} else {
+		struct edgetpu_usage_header *header = buf;
+
+		metric_size = header->metric_size;
+		num_metrics = header->num_metrics;
+		version = header->version;
+		metric = (struct edgetpu_usage_metric *)((char *)header + header->header_bytes);
+	}
+
+	etdev_dbg(etdev, "%s: v=%u n=%u sz=%u", __func__, version, num_metrics, metric_size);
+	if (metric_size < EDGETPU_USAGE_METRIC_SIZE_V1) {
 		etdev_warn_once(etdev, "fw metric size %u less than minimum %u",
-				header->metric_size, EDGETPU_USAGE_METRIC_SIZE_V1);
+				metric_size, EDGETPU_USAGE_METRIC_SIZE_V1);
 		return;
 	}
 
-	if (header->metric_size > sizeof(struct edgetpu_usage_metric))
+	if (metric_size > sizeof(struct edgetpu_usage_metric))
 		etdev_dbg(etdev, "fw metrics are later version with unknown fields");
 
-	for (i = 0; i < header->num_metrics; i++) {
+	for (i = 0; i < num_metrics; i++) {
 		switch (metric->type) {
 		case EDGETPU_METRIC_TYPE_TPU_USAGE:
 			edgetpu_usage_add(etdev, &metric->tpu_usage);
@@ -287,7 +308,7 @@ void edgetpu_usage_stats_process_buffer(struct edgetpu_dev *etdev, void *buf)
 			break;
 		}
 
-		metric = (struct edgetpu_usage_metric *)((char *)metric + header->metric_size);
+		metric = (struct edgetpu_usage_metric *)((char *)metric + metric_size);
 	}
 }
 
@@ -915,6 +936,7 @@ static struct attribute *usage_stats_dev_attrs[] = {
 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;
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 6a89acb..ee908e1 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -11,7 +11,7 @@
 #include <linux/mutex.h>
 
 /* The highest version of usage metrics handled by this driver. */
-#define EDGETPU_USAGE_METRIC_VERSION	1
+#define EDGETPU_USAGE_METRIC_VERSION	2
 
 /*
  * Size in bytes of usage metric v1.
@@ -21,9 +21,17 @@
  */
 #define EDGETPU_USAGE_METRIC_SIZE_V1	20
 
+/* v1 metric header struct. */
+struct edgetpu_usage_header_v1 {
+	uint32_t num_metrics;		/* Number of metrics being reported */
+	uint32_t metric_size;		/* Size of each metric struct */
+};
+
 /* Header struct in the metric buffer. */
 /* Must be kept in sync with firmware struct UsageTrackerHeader */
 struct edgetpu_usage_header {
+	uint16_t header_bytes;		/* Number of bytes in this header */
+	uint16_t version;		/* Metrics version */
 	uint32_t num_metrics;		/* Number of metrics being reported */
 	uint32_t metric_size;		/* Size of each metric struct */
 };
@@ -31,15 +39,24 @@ struct edgetpu_usage_header {
 /*
  * Encapsulate TPU core usage information of a specific application for a
  * specific power state.
- * Must be kept in sync with firmware struct TpuUsage.
+ * Must be kept in sync with firmware struct CoreUsage.
  */
 struct tpu_usage {
 	/* Unique identifier of the application. */
 	int32_t uid;
 	/* The power state of the device (values are chip dependent) */
+	/* Now called operating_point in FW. */
 	uint32_t power_state;
 	/* Duration of usage in microseconds. */
 	uint32_t duration_us;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Compute Core: TPU cluster ID. */
+	/* Called core_id in FW. */
+	uint8_t cluster_id;
+	/* Reserved.  Filling out the next 32-bit boundary. */
+	uint8_t reserved[3];
 };
 
 /*
@@ -49,9 +66,12 @@ struct tpu_usage {
 enum edgetpu_usage_component {
 	/* The device as a whole */
 	EDGETPU_USAGE_COMPONENT_DEVICE = 0,
-	/* Just the TPU core */
+	/* Just the TPU core (scalar core and tiles) */
 	EDGETPU_USAGE_COMPONENT_TPU = 1,
-	EDGETPU_USAGE_COMPONENT_COUNT = 2, /* number of components above */
+	/* Control core (ARM Cortex-R52 CPU) */
+	EDGETPU_USAGE_COMPONENT_CONTROLCORE = 2,
+
+	EDGETPU_USAGE_COMPONENT_COUNT = 3, /* number of components above */
 };
 
 /*
@@ -73,7 +93,7 @@ enum edgetpu_usage_counter_type {
 	EDGETPU_COUNTER_TPU_ACTIVE_CYCLES = 0,
 	/* Number of stalls caused by throttling. */
 	EDGETPU_COUNTER_TPU_THROTTLE_STALLS = 1,
-	/* Number of graph invocations. */
+	/* Number of graph invocations. (Now called kWorkload in FW.) */
 	EDGETPU_COUNTER_INFERENCES = 2,
 	/* Number of TPU offload op invocations. */
 	EDGETPU_COUNTER_TPU_OPS = 3,
@@ -92,7 +112,12 @@ enum edgetpu_usage_counter_type {
 	/* Number of times (firmware)suspend function takes longer than SLA time. */
 	EDGETPU_COUNTER_LONG_SUSPEND = 10,
 
-	EDGETPU_COUNTER_COUNT = 11, /* number of counters above */
+	/* The following counters are added in metrics v2. */
+
+	/* Number of context switches on a compute core. */
+	EDGETPU_COUNTER_CONTEXT_SWITCHES = 11,
+
+	EDGETPU_COUNTER_COUNT = 12, /* number of counters above */
 };
 
 /* Generic counter. Only reported if it has a value larger than 0. */
@@ -102,6 +127,11 @@ struct __packed edgetpu_usage_counter {
 
 	/* Accumulated value since last initialization. */
 	uint64_t value;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Reporting component. */
+	uint8_t component_id;
 };
 
 /* Defines different max watermarks we track. */
@@ -132,15 +162,21 @@ struct __packed edgetpu_usage_max_watermark {
 	 * non-mobile, firmware boot on mobile).
 	 */
 	uint64_t value;
+
+	/* Following fields are added in metrics v2 */
+
+	/* Reporting component. */
+	uint8_t component_id;
 };
 
 /* An enum to identify the tracked firmware threads. */
 /* Must be kept in sync with firmware enum class UsageTrackerThreadId. */
 enum edgetpu_usage_threadid {
-	/* Individual thread IDs are not tracked. */
+	/* Individual thread IDs do not have identifiers assigned. */
+	/* Thread ID 14, used for other IP, is not used for TPU */
 
 	/* Number of task identifiers. */
-	EDGETPU_FW_THREAD_COUNT = 12,
+	EDGETPU_FW_THREAD_COUNT = 14,
 };
 
 /* Statistics related to a single thread in firmware. */
@@ -184,6 +220,8 @@ struct edgetpu_usage_metric {
 #define UID_HASH_BITS 3
 
 struct edgetpu_usage_stats {
+	/* if true the current firmware only implements metrics V1 */
+	bool use_metrics_v1;
 	DECLARE_HASHTABLE(uid_hash_table, UID_HASH_BITS);
 	/* component utilization values reported by firmware */
 	int32_t component_utilization[EDGETPU_USAGE_COMPONENT_COUNT];
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index 7de0874..7af6c7e 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -11,6 +11,7 @@ gcip-objs		:= gcip-alloc-helper.o \
 			   gcip-domain-pool.o \
 			   gcip-firmware.o \
 			   gcip-image-config.o \
+			   gcip-iommu.o \
 			   gcip-kci.o \
 			   gcip-mailbox.o \
 			   gcip-mem-pool.o \
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c
new file mode 100644
index 0000000..2e0dac6
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-iommu.c
@@ -0,0 +1,432 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Manages GCIP IOMMU domains and allocates/maps IOVAs.
+ *
+ * Copyright (C) 2023 Google LLC
+ */
+
+#include <linux/bitops.h>
+#include <linux/device.h>
+#include <linux/dma-direction.h>
+#include <linux/dma-iommu.h>
+#include <linux/dma-mapping.h>
+#include <linux/genalloc.h>
+#include <linux/iova.h>
+#include <linux/log2.h>
+#include <linux/of.h>
+#include <linux/scatterlist.h>
+#include <linux/slab.h>
+
+#include <gcip/gcip-domain-pool.h>
+#include <gcip/gcip-iommu.h>
+#include <gcip/gcip-mem-pool.h>
+
+#define HAS_IOVAD_BEST_FIT_ALGO (IS_ENABLED(CONFIG_GCIP_TEST) || IS_ENABLED(CONFIG_ANDROID))
+
+/* Macros for manipulating @gcip_map_flags parameter. */
+#define GCIP_MAP_FLAGS_GET_VALUE(ATTR, flags)                                                      \
+	(((flags) >> (GCIP_MAP_FLAGS_##ATTR##_OFFSET)) &                                           \
+	 (BIT_ULL(GCIP_MAP_FLAGS_##ATTR##_BIT_SIZE) - 1))
+#define GCIP_MAP_FLAGS_GET_DMA_DIRECTION(flags) GCIP_MAP_FLAGS_GET_VALUE(DMA_DIRECTION, flags)
+#define GCIP_MAP_FLAGS_GET_DMA_COHERENT(flags) GCIP_MAP_FLAGS_GET_VALUE(DMA_COHERENT, flags)
+#define GCIP_MAP_FLAGS_GET_DMA_ATTR(flags) GCIP_MAP_FLAGS_GET_VALUE(DMA_ATTR, flags)
+
+/**
+ * dma_info_to_prot - Translate DMA API directions and attributes to IOMMU API
+ *                    page flags.
+ * @dir: Direction of DMA transfer
+ * @coherent: If true, create coherent mappings of the scatterlist.
+ * @attrs: DMA attributes for the mapping
+ *
+ * See v5.15.94/source/drivers/iommu/dma-iommu.c#L418
+ *
+ * Return: corresponding IOMMU API page protection flags
+ */
+static int dma_info_to_prot(enum dma_data_direction dir, bool coherent, unsigned long attrs)
+{
+	int prot = coherent ? IOMMU_CACHE : 0;
+
+	if (attrs & DMA_ATTR_PRIVILEGED)
+		prot |= IOMMU_PRIV;
+
+	switch (dir) {
+	case DMA_BIDIRECTIONAL:
+		return prot | IOMMU_READ | IOMMU_WRITE;
+	case DMA_TO_DEVICE:
+		return prot | IOMMU_READ;
+	case DMA_FROM_DEVICE:
+		return prot | IOMMU_WRITE;
+	default:
+		return 0;
+	}
+}
+
+static inline unsigned long gcip_iommu_domain_shift(struct gcip_iommu_domain *domain)
+{
+	return __ffs(domain->domain_pool->granule);
+}
+
+static inline unsigned long gcip_iommu_domain_pfn(struct gcip_iommu_domain *domain, dma_addr_t iova)
+{
+	return iova >> gcip_iommu_domain_shift(domain);
+}
+
+static inline size_t gcip_iommu_domain_align(struct gcip_iommu_domain *domain, size_t size)
+{
+	return ALIGN(size, domain->domain_pool->granule);
+}
+
+static int iovad_initialize_domain(struct gcip_iommu_domain *domain)
+{
+	struct gcip_iommu_domain_pool *dpool = domain->domain_pool;
+
+	init_iova_domain(&domain->iova_space.iovad, dpool->granule,
+			 max_t(unsigned long, 1, dpool->base_daddr >> ilog2(dpool->granule)));
+
+	return 0;
+}
+
+static void iovad_finalize_domain(struct gcip_iommu_domain *domain)
+{
+	put_iova_domain(&domain->iova_space.iovad);
+}
+
+static void iovad_enable_best_fit_algo(struct gcip_iommu_domain *domain)
+{
+#if HAS_IOVAD_BEST_FIT_ALGO
+	domain->iova_space.iovad.best_fit = true;
+#endif /* HAS_IOVAD_BEST_FIT_ALGO */
+}
+
+static dma_addr_t iovad_alloc_iova_space(struct gcip_iommu_domain *domain, size_t size)
+{
+	unsigned long iova, shift = gcip_iommu_domain_shift(domain);
+
+	iova = alloc_iova_fast(&domain->iova_space.iovad, size >> shift,
+			       domain->domain_pool->last_daddr >> shift, true);
+
+	return (dma_addr_t)iova << shift;
+}
+
+static void iovad_free_iova_space(struct gcip_iommu_domain *domain, dma_addr_t iova, size_t size)
+{
+	free_iova_fast(&domain->iova_space.iovad, gcip_iommu_domain_pfn(domain, iova),
+		       size >> gcip_iommu_domain_shift(domain));
+}
+
+static const struct gcip_iommu_domain_ops iovad_ops = {
+	.initialize_domain = iovad_initialize_domain,
+	.finalize_domain = iovad_finalize_domain,
+	.enable_best_fit_algo = iovad_enable_best_fit_algo,
+	.alloc_iova_space = iovad_alloc_iova_space,
+	.free_iova_space = iovad_free_iova_space,
+};
+
+static int mem_pool_initialize_domain(struct gcip_iommu_domain *domain)
+{
+	struct gcip_iommu_domain_pool *dpool = domain->domain_pool;
+	int ret;
+
+	ret = gcip_mem_pool_init(&domain->iova_space.mem_pool, dpool->dev, dpool->base_daddr,
+				 dpool->size, dpool->granule);
+
+	return ret;
+}
+
+static void mem_pool_finalize_domain(struct gcip_iommu_domain *domain)
+{
+	gcip_mem_pool_exit(&domain->iova_space.mem_pool);
+}
+
+static void mem_pool_enable_best_fit_algo(struct gcip_iommu_domain *domain)
+{
+	gen_pool_set_algo(domain->iova_space.mem_pool.gen_pool, gen_pool_best_fit, NULL);
+}
+
+static dma_addr_t mem_pool_alloc_iova_space(struct gcip_iommu_domain *domain, size_t size)
+{
+	return (dma_addr_t)gcip_mem_pool_alloc(&domain->iova_space.mem_pool, size);
+}
+
+static void mem_pool_free_iova_space(struct gcip_iommu_domain *domain, dma_addr_t iova, size_t size)
+{
+	gcip_mem_pool_free(&domain->iova_space.mem_pool, iova, size);
+}
+
+static const struct gcip_iommu_domain_ops mem_pool_ops = {
+	.initialize_domain = mem_pool_initialize_domain,
+	.finalize_domain = mem_pool_finalize_domain,
+	.enable_best_fit_algo = mem_pool_enable_best_fit_algo,
+	.alloc_iova_space = mem_pool_alloc_iova_space,
+	.free_iova_space = mem_pool_free_iova_space,
+};
+
+static bool enable_best_fit_algo_legacy(struct gcip_iommu_domain_pool *pool)
+{
+	__maybe_unused int ret;
+
+#if HAS_IOVAD_BEST_FIT_ALGO
+	ret = iommu_dma_enable_best_fit_algo(pool->dev);
+	if (!ret)
+		return true;
+	dev_warn(pool->dev, "Failed to enable best-fit IOMMU domain pool (%d)\n", ret);
+#else
+	dev_warn(pool->dev, "This env doesn't support best-fit algorithm in the legacy mode");
+#endif
+	return false;
+}
+
+static ssize_t dma_iommu_map_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl,
+				int nents, enum dma_data_direction dir, unsigned long attrs,
+				int prot)
+{
+	int nents_mapped;
+	dma_addr_t iova;
+	ssize_t ret;
+
+	nents_mapped = dma_map_sg_attrs(domain->dev, sgl, nents, dir, attrs);
+	if (!nents_mapped)
+		return 0;
+
+	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;
+	}
+
+	return nents_mapped;
+}
+
+static void dma_iommu_unmap_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl, int nents,
+			       enum dma_data_direction dir, unsigned long attrs)
+{
+	struct scatterlist *sg;
+	size_t size = 0;
+	int i;
+
+	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");
+	dma_unmap_sg_attrs(domain->dev, sgl, nents, dir, attrs);
+}
+
+int gcip_iommu_domain_pool_init(struct gcip_iommu_domain_pool *pool, struct device *dev,
+				dma_addr_t base_daddr, size_t iova_space_size, size_t granule,
+				unsigned int num_domains, enum gcip_iommu_domain_type domain_type)
+{
+	const __be32 *user_window;
+	int ret;
+
+	ret = gcip_domain_pool_init(dev, &pool->domain_pool, num_domains);
+	if (ret)
+		return ret;
+
+	pool->dev = dev;
+	pool->base_daddr = base_daddr;
+	pool->size = iova_space_size;
+	pool->granule = granule;
+	pool->best_fit = false;
+	pool->domain_type = domain_type;
+
+	if (!base_daddr || !iova_space_size) {
+		user_window = of_get_property(dev->of_node, "gcip-dma-window", NULL);
+		if (!user_window) {
+			dev_warn(dev, "Failed to find gcip-dma-window property");
+		} else {
+			pool->base_daddr = of_read_number(user_window, 1);
+			pool->size = of_read_number(user_window + 1, 1);
+		}
+	}
+
+	if (!pool->base_daddr || !pool->size) {
+		dev_warn(dev, "GCIP IOMMU domain pool is initialized as the legacy mode");
+		pool->size = 0;
+	} else {
+		pool->last_daddr = pool->base_daddr + pool->size - 1;
+	}
+
+	dev_dbg(dev, "Init GCIP IOMMU domain pool, base_daddr=%#llx, size=%#zx", pool->base_daddr,
+		pool->size);
+
+	return 0;
+}
+
+void gcip_iommu_domain_pool_destroy(struct gcip_iommu_domain_pool *pool)
+{
+	gcip_domain_pool_destroy(&pool->domain_pool);
+}
+
+void gcip_iommu_domain_pool_enable_best_fit_algo(struct gcip_iommu_domain_pool *pool)
+{
+	if (gcip_iommu_domain_pool_is_legacy_mode(pool)) {
+		pool->best_fit = enable_best_fit_algo_legacy(pool);
+	} else if (pool->domain_type == GCIP_IOMMU_DOMAIN_TYPE_IOVAD && !HAS_IOVAD_BEST_FIT_ALGO) {
+		dev_warn(pool->dev, "This env doesn't support best-fit algorithm with IOVAD");
+		pool->best_fit = false;
+	} else {
+		pool->best_fit = true;
+	}
+}
+
+void gcip_iommu_domain_pool_enable_legacy_mode(struct gcip_iommu_domain_pool *pool)
+{
+	pool->size = 0;
+	pool->base_daddr = 0;
+
+	if (pool->best_fit)
+		pool->best_fit = enable_best_fit_algo_legacy(pool);
+}
+
+struct gcip_iommu_domain *gcip_iommu_domain_pool_alloc_domain(struct gcip_iommu_domain_pool *pool)
+{
+	struct gcip_iommu_domain *gdomain;
+	int ret;
+
+	gdomain = devm_kzalloc(pool->dev, sizeof(*gdomain), GFP_KERNEL);
+	if (!gdomain)
+		return ERR_PTR(-ENOMEM);
+
+	gdomain->dev = pool->dev;
+	gdomain->domain_pool = pool;
+	gdomain->domain = gcip_domain_pool_alloc(&pool->domain_pool);
+	if (IS_ERR_OR_NULL(gdomain->domain)) {
+		ret = -ENOMEM;
+		goto err_free_gdomain;
+	}
+
+	if (gcip_iommu_domain_pool_is_legacy_mode(pool)) {
+		gdomain->legacy_mode = true;
+		return gdomain;
+	}
+
+	switch (pool->domain_type) {
+	case GCIP_IOMMU_DOMAIN_TYPE_IOVAD:
+		gdomain->ops = &iovad_ops;
+		break;
+	case GCIP_IOMMU_DOMAIN_TYPE_MEM_POOL:
+		gdomain->ops = &mem_pool_ops;
+		break;
+	default:
+		ret = -EINVAL;
+		goto err_free_domain_pool;
+	}
+
+	ret = gdomain->ops->initialize_domain(gdomain);
+	if (ret)
+		goto err_free_domain_pool;
+
+	if (pool->best_fit)
+		gdomain->ops->enable_best_fit_algo(gdomain);
+
+	return gdomain;
+
+err_free_domain_pool:
+	gcip_domain_pool_free(&pool->domain_pool, gdomain->domain);
+err_free_gdomain:
+	devm_kfree(pool->dev, gdomain);
+	return ERR_PTR(ret);
+}
+
+void gcip_iommu_domain_pool_free_domain(struct gcip_iommu_domain_pool *pool,
+					struct gcip_iommu_domain *domain)
+{
+	if (!gcip_iommu_domain_is_legacy_mode(domain))
+		domain->ops->finalize_domain(domain);
+	gcip_domain_pool_free(&pool->domain_pool, domain->domain);
+	devm_kfree(pool->dev, domain);
+}
+
+unsigned int gcip_iommu_domain_map_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl,
+				      int nents, u64 gcip_map_flags)
+{
+	enum dma_data_direction dir = GCIP_MAP_FLAGS_GET_DMA_DIRECTION(gcip_map_flags);
+	bool coherent = GCIP_MAP_FLAGS_GET_DMA_COHERENT(gcip_map_flags);
+	unsigned long attrs = GCIP_MAP_FLAGS_GET_DMA_ATTR(gcip_map_flags);
+	int i, prot = dma_info_to_prot(dir, coherent, attrs);
+	struct scatterlist *sg;
+	dma_addr_t iova;
+	size_t iova_len = 0;
+	ssize_t ret;
+
+	if (gcip_iommu_domain_is_legacy_mode(domain))
+		return dma_iommu_map_sg(domain, sgl, nents, dir, attrs, prot);
+
+	/* Calculates how much IOVA space we need. */
+	for_each_sg (sgl, sg, nents, i)
+		iova_len += sg->length;
+
+	/* Allocates one continuous IOVA. */
+	iova = domain->ops->alloc_iova_space(domain, gcip_iommu_domain_align(domain, iova_len));
+	if (!iova)
+		return 0;
+
+	/*
+	 * Maps scatterlist to the allocated IOVA.
+	 *
+	 * It will iterate each scatter list segment in order and map them to the IOMMU domain
+	 * as amount of the size of each segment successively.
+	 * Returns an error on failure or the total length of mapped segments on success.
+	 *
+	 * Note: Before Linux 5.15, its return type was `size_t` and it returned 0 on failure.
+	 *       To make it compatible with those old versions, we should cast the return value.
+	 */
+	ret = (ssize_t)iommu_map_sg(domain->domain, iova, sgl, nents, prot);
+	if (ret < 0 || ret < iova_len)
+		goto err_free_iova;
+
+	/* Fills out the mapping information. */
+	sg_dma_address(sgl) = iova;
+	sg_dma_len(sgl) = iova_len;
+
+	/* As it put the whole mapping information to the first segment, it should return 1. */
+	return 1;
+
+err_free_iova:
+	domain->ops->free_iova_space(domain, iova, gcip_iommu_domain_align(domain, iova_len));
+	return 0;
+}
+
+void gcip_iommu_domain_unmap_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl,
+				int nents, u64 gcip_map_flags)
+{
+	dma_addr_t iova;
+	size_t iova_len;
+
+	if (gcip_iommu_domain_is_legacy_mode(domain)) {
+		enum dma_data_direction dir = GCIP_MAP_FLAGS_GET_DMA_DIRECTION(gcip_map_flags);
+		unsigned long attrs = GCIP_MAP_FLAGS_GET_DMA_ATTR(gcip_map_flags);
+
+		dma_iommu_unmap_sg(domain, sgl, nents, dir, attrs);
+		return;
+	}
+
+	iova = sg_dma_address(sgl);
+	iova_len = sg_dma_len(sgl);
+
+	iommu_unmap(domain->domain, iova, iova_len);
+	domain->ops->free_iova_space(domain, iova, gcip_iommu_domain_align(domain, iova_len));
+}
+
+struct gcip_iommu_domain *gcip_iommu_get_domain_for_dev(struct device *dev)
+{
+	struct gcip_iommu_domain *gdomain;
+
+	gdomain = devm_kzalloc(dev, sizeof(*gdomain), GFP_KERNEL);
+	if (!gdomain)
+		return ERR_PTR(-ENOMEM);
+
+	gdomain->domain = iommu_get_domain_for_dev(dev);
+	if (!gdomain->domain) {
+		devm_kfree(dev, gdomain);
+		return ERR_PTR(-ENODEV);
+	}
+
+	gdomain->dev = dev;
+	gdomain->legacy_mode = true;
+
+	return gdomain;
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h
new file mode 100644
index 0000000..4e04b7e
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-iommu.h
@@ -0,0 +1,267 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+/*
+ * Manages GCIP IOMMU domains and allocates/maps IOVAs.
+ *
+ * One can replace allocating IOVAs via Linux DMA interface which will allocate and map them to
+ * the default IOMMU domain with this framework. This framework will allocate and map IOVAs to the
+ * specific IOMMU domain directly. This has following two advantages:
+ *
+ * - Can remove the mapping time by once as it maps to the target IOMMU domain directly.
+ * - IOMMU domains don't have to share the total capacity.
+ *
+ * GCIP IOMMU domain is implemented by utilizing multiple kinds of IOVA space pool:
+ * - struct iova_domain
+ * - struct gcip_mem_pool
+ *
+ * Copyright (C) 2023 Google LLC
+ */
+
+#ifndef __GCIP_IOMMU_H__
+#define __GCIP_IOMMU_H__
+
+#include <linux/device.h>
+#include <linux/iommu.h>
+#include <linux/iova.h>
+#include <linux/scatterlist.h>
+
+#include <gcip/gcip-domain-pool.h>
+#include <gcip/gcip-mem-pool.h>
+
+/*
+ * Helpers for manipulating @gcip_map_flags parameter of the `gcip_iommu_domain_{map,unmap}_sg`
+ * functions.
+ */
+#define GCIP_MAP_FLAGS_DMA_DIRECTION_OFFSET 0
+#define GCIP_MAP_FLAGS_DMA_DIRECTION_BIT_SIZE 2
+#define GCIP_MAP_FLAGS_DMA_DIRECTION_TO_FLAGS(dir)                                                 \
+	((u64)(dir) << GCIP_MAP_FLAGS_DMA_DIRECTION_OFFSET)
+
+#define GCIP_MAP_FLAGS_DMA_COHERENT_OFFSET                                                         \
+	(GCIP_MAP_FLAGS_DMA_DIRECTION_OFFSET + GCIP_MAP_FLAGS_DMA_DIRECTION_BIT_SIZE)
+#define GCIP_MAP_FLAGS_DMA_COHERENT_BIT_SIZE 1
+#define GCIP_MAP_FLAGS_DMA_COHERENT_TO_FLAGS(coherent)                                             \
+	((u64)(coherent) << GCIP_MAP_FLAGS_DMA_COHERENT_OFFSET)
+
+#define GCIP_MAP_FLAGS_DMA_ATTR_OFFSET                                                             \
+	(GCIP_MAP_FLAGS_DMA_COHERENT_OFFSET + GCIP_MAP_FLAGS_DMA_COHERENT_BIT_SIZE)
+#define GCIP_MAP_FLAGS_DMA_ATTR_BIT_SIZE 10
+#define GCIP_MAP_FLAGS_DMA_ATTR_TO_FLAGS(attr) ((u64)(attr) << GCIP_MAP_FLAGS_DMA_ATTR_OFFSET)
+
+struct gcip_iommu_domain_ops;
+
+/*
+ * Type of IOVA space pool that IOMMU domain will utilize.
+ * Regardless of the type, its functionality will be the same. However, its implementation might be
+ * different. For example, iova_domain uses red-black tree for the memory management, but gen_pool
+ * uses bitmap. Therefore, their performance might be different and the kernel drivers can choose
+ * which one to use according to its real use cases and the performance.
+ *
+ * Note: in legacy mode, only iova_domain is available as the Linux implementation utilizes that.
+ */
+enum gcip_iommu_domain_type {
+	/* Uses iova_domain. */
+	GCIP_IOMMU_DOMAIN_TYPE_IOVAD,
+	/* Uses gcip_mem_pool which is based on gen_pool. */
+	GCIP_IOMMU_DOMAIN_TYPE_MEM_POOL,
+};
+
+/*
+ * IOMMU domain pool.
+ *
+ * It manages the pool of IOMMU domains. Also, it specifies the base address and the size of IOMMU
+ * domains. Also, one can choose the data structure and algorithm of IOVA space management.
+ */
+struct gcip_iommu_domain_pool {
+	struct device *dev;
+	struct gcip_domain_pool domain_pool;
+	dma_addr_t base_daddr;
+	/* Will hold (base_daddr + size - 1) to prevent calculating it every IOVAD mappings. */
+	dma_addr_t last_daddr;
+	size_t size;
+	size_t granule;
+	bool best_fit;
+	enum gcip_iommu_domain_type domain_type;
+};
+
+/*
+ * Wrapper of iommu_domain.
+ * It has its own IOVA space pool based on iova_domain or gcip_mem_pool. One can choose one of them
+ * when calling the `gcip_iommu_domain_pool_init` function. See `enum gcip_iommu_domain_type`
+ * for details.
+ */
+struct gcip_iommu_domain {
+	struct device *dev;
+	struct gcip_iommu_domain_pool *domain_pool;
+	struct iommu_domain *domain;
+	bool legacy_mode;
+	union {
+		struct iova_domain iovad;
+		struct gcip_mem_pool mem_pool;
+	} iova_space;
+	const struct gcip_iommu_domain_ops *ops;
+};
+
+/*
+ * Holds operators which will be set according to the @domain_type.
+ * These callbacks will be filled automatically when a `struct gcip_iommu_domain` is allocated.
+ */
+struct gcip_iommu_domain_ops {
+	/* Initializes pool of @domain. */
+	int (*initialize_domain)(struct gcip_iommu_domain *domain);
+	/* Destroyes pool of @domain */
+	void (*finalize_domain)(struct gcip_iommu_domain *domain);
+	/*
+	 * Enables best-fit algorithm for the memory management.
+	 * Only domains which are allocated after calling this callback will be affected.
+	 */
+	void (*enable_best_fit_algo)(struct gcip_iommu_domain *domain);
+	/* Allocates @size of buffer and returns its IOVA. */
+	dma_addr_t (*alloc_iova_space)(struct gcip_iommu_domain *domain, size_t size);
+	/* Releases @size of buffer which was allocated to @iova. */
+	void (*free_iova_space)(struct gcip_iommu_domain *domain, dma_addr_t iova, size_t size);
+};
+
+/*
+ * Initializes an IOMMU domain pool.
+ *
+ * One can specify the base DMA address and IOVA space size via @base_daddr and @iova_space_size
+ * parameters. If any of them is 0, it will try to parse "gcip-dma-window" property from the device
+ * tree of @dev.
+ *
+ * If the base DMA address and IOVA space size are set successfully (i.e., larger than 0), IOMMU
+ * domains allocated by this domain pool will have their own IOVA space pool and will map buffers
+ * to their own IOMMU domain directly.
+ *
+ * Otherwise, it will fall into the legacy mode which will utilize the native DMA-IOMMU APIs.
+ * In this mode, it will map the buffer to the default IOMMU domain first and then remap it to the
+ * target domain.
+ *
+ * @pool: IOMMU domain pool to be initialized.
+ * @dev: Device where to parse "gcip-dma-window" property.
+ * @base_addr: The base address of IOVA space. Must be greater than 0 and a multiple of @granule.
+ * @iova_space_size: The size of the IOVA space. @size must be a multiple of @granule.
+ * @granule: The granule when invoking the IOMMU domain pool. Must be a power of 2.
+ * @num_domains: The number of IOMMU domains.
+ * @domain_type: Type of the IOMMU domain.
+ *
+ * Returns 0 on success or negative error value.
+ */
+int gcip_iommu_domain_pool_init(struct gcip_iommu_domain_pool *pool, struct device *dev,
+				dma_addr_t base_daddr, size_t iova_space_size, size_t granule,
+				unsigned int num_domains, enum gcip_iommu_domain_type domain_type);
+
+/*
+ * Destroys an IOMMU domain pool.
+ *
+ * @pool: IOMMU domain pool to be destroyed.
+ */
+void gcip_iommu_domain_pool_destroy(struct gcip_iommu_domain_pool *pool);
+
+/*
+ * Enables the best fit algorithm for allocating an IOVA space.
+ * It affects domains which are allocated after calling this function only.
+ *
+ * @pool: IOMMU domain pool to be enabled.
+ */
+void gcip_iommu_domain_pool_enable_best_fit_algo(struct gcip_iommu_domain_pool *pool);
+
+/*
+ * Enables the legacy mode of allocating and mapping IOVA logic which utilizes native DMA-IOMMU
+ * APIs of the Linux kernel.
+ * It affects domains which are allocated after calling this function only.
+ *
+ * @pool: IOMMU domain pool to be enabled.
+ */
+void gcip_iommu_domain_pool_enable_legacy_mode(struct gcip_iommu_domain_pool *pool);
+
+/*
+ * Returns whether @pool is using legacy mode or not.
+ *
+ * @pool: IOMMU domain pool to be checked.
+ */
+static inline bool gcip_iommu_domain_pool_is_legacy_mode(struct gcip_iommu_domain_pool *pool)
+{
+	return !(pool && pool->size);
+}
+
+/*
+ * Allocates a GCIP IOMMU domain.
+ *
+ * @pool: IOMMU domain pool.
+ *
+ * Returns a pointer of allocated domain on success or an error pointer on failure.
+ */
+struct gcip_iommu_domain *gcip_iommu_domain_pool_alloc_domain(struct gcip_iommu_domain_pool *pool);
+
+/*
+ * Releases a GCIP IOMMU domain.
+ *
+ * Before calling this function, you must unmap all IOVAs by calling `gcip_iommu_domain_unmap{_sg}`
+ * functions.
+ *
+ * @pool: IOMMU domain pool.
+ * @domain: GCIP IOMMU domain to be released.
+ */
+void gcip_iommu_domain_pool_free_domain(struct gcip_iommu_domain_pool *pool,
+					struct gcip_iommu_domain *domain);
+
+/*
+ * Returns whether @domain is using legacy mode or not.
+ *
+ * @domain: GCIP IOMMU domain to be checked.
+ */
+static inline bool gcip_iommu_domain_is_legacy_mode(struct gcip_iommu_domain *domain)
+{
+	return domain->legacy_mode;
+}
+
+/*
+ * Allocates an IOVA for the scatterlist and maps it to @domain.
+ *
+ * @domain: GCIP IOMMU domain which manages IOVA addresses.
+ * @sgl: Scatterlist to be mapped.
+ * @nents: The number of entries in @sgl.
+ * @gcip_map_flags: Flags indicating mapping attributes.
+ *
+ * Bitfields:
+ *   [1:0]   - DMA_DIRECTION:
+ *               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)
+ *   [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)
+ *   [63:13] - RESERVED
+ *               Set RESERVED bits to 0 to ensure backwards compatibility.
+ *
+ * One can use `GCIP_MAP_FLAGS_DMA_*_TO_FLAGS` macros to generate a flag.
+ *
+ * Returns the number of entries which are mapped to @domain. Returns 0 if it fails.
+ */
+unsigned int gcip_iommu_domain_map_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl,
+				      int nents, u64 gcip_map_flags);
+
+/*
+ * Unmaps an IOVA which was mapped for the scatterlist.
+ *
+ * @domain: GCIP IOMMU domain which manages IOVA addresses.
+ * @sgl: Scatterlist to be unmapped.
+ * @gcip_map_flags: The same as the `gcip_iommu_domain_map_sg` function.
+ *                  It will be ignored in the non-legacy mode.
+ */
+void gcip_iommu_domain_unmap_sg(struct gcip_iommu_domain *domain, struct scatterlist *sgl,
+				int nents, u64 gcip_map_flags);
+
+/*
+ * Returns a default GCIP IOMMU domain.
+ * This domain works with the legacy mode only.
+ *
+ * @dev: Device where to fetch the default IOMMU domain.
+ */
+struct gcip_iommu_domain *gcip_iommu_get_domain_for_dev(struct device *dev);
+
+#endif /* __GCIP_IOMMU_H__ */
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 eb83550..2aa721b 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
@@ -90,6 +90,9 @@ enum gcip_kci_code {
 	GCIP_KCI_CODE_OPEN_DEVICE = 9,
 	GCIP_KCI_CODE_CLOSE_DEVICE = 10,
 	GCIP_KCI_CODE_FIRMWARE_INFO = 11,
+	/* TODO(b/271372136): remove v1 when v1 firmware no longer in use. */
+	GCIP_KCI_CODE_GET_USAGE_V1 = 12,
+	/* Backward compatible define, also update when v1 firmware no longer in use. */
 	GCIP_KCI_CODE_GET_USAGE = 12,
 	GCIP_KCI_CODE_NOTIFY_THROTTLING = 13,
 	GCIP_KCI_CODE_BLOCK_BUS_SPEED_CONTROL = 14,
@@ -99,6 +102,7 @@ enum gcip_kci_code {
 	GCIP_KCI_CODE_UNLINK_OFFLOAD_VMBOX = 18,
 	GCIP_KCI_CODE_FIRMWARE_TRACING_LEVEL = 19,
 	GCIP_KCI_CODE_THERMAL_CONTROL = 20,
+	GCIP_KCI_CODE_GET_USAGE_V2 = 21,
 
 	GCIP_KCI_CODE_RKCI_ACK = 256,
 };
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 4842598..1e6ce05 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-pm.h
@@ -40,12 +40,14 @@ struct gcip_pm_args {
 	void *data;
 	/*
 	 * Device-specific power up.
-	 * Called with @pm->lock hold and nesting is handled at generic layer.
+	 * Called with @pm->lock held and nesting is handled at generic layer.
+	 * The IP driver may reject power on for such conditions as thermal suspend in this
+	 * callback.
 	 */
 	int (*power_up)(void *data);
 	/*
 	 * Device-specific power down.
-	 * Called with @pm->lock hold and nesting is handled at generic layer.
+	 * Called with @pm->lock held and nesting is handled at generic layer.
 	 * Returning -EAGAIN will trigger a retry after GCIP_ASYNC_POWER_DOWN_RETRY_DELAY ms.
 	 */
 	int (*power_down)(void *data);
@@ -106,7 +108,7 @@ bool gcip_pm_is_powered(struct gcip_pm *pm);
 /* Shuts down the device if @pm->count equals to 0 or @force is true. */
 void gcip_pm_shutdown(struct gcip_pm *pm, bool force);
 
-/* Make sure @pm->lock is hold. */
+/* Make sure @pm->lock is held. */
 static inline void gcip_pm_lockdep_assert_held(struct gcip_pm *pm)
 {
 	if (!pm)
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 2c9eb1b..53571e0 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -11,6 +11,8 @@
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
 
+#include <gcip/gcip-thermal.h>
+
 #include "edgetpu-config.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
@@ -196,6 +198,12 @@ static int mobile_power_up(void *data)
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 	int ret;
 
+	if (gcip_thermal_is_device_suspended(etdev->thermal)) {
+		etdev_warn_ratelimited(etdev,
+				       "power up rejected due to device thermal limit exceeded");
+		return -EAGAIN;
+	}
+
 	if (platform_pwr->is_block_down) {
 		int times = 0;