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

edgetpu: Kbuild: Add back EXTRA_CFLAGS
Bug: 265616915
gcip: expose config size calculation to header
Bug: 265605775
gcip: expose config size calculation to header
Bug: 265605775 (repeat)
gcip: Use strscpy instead of memcpy
gcip: enhance image config NS mapping decoding
Bug: 265565307
gcip: sync RKCI codes with fw
gcip: update comments of gcip_dma_fence_init
Bug: 264220687
gcip: add to_gcip_fence to gcip-dma-fence
Bug: 264220687 (repeat)
gcip: fix GCIP_DMA_FENCE_LIST_UNLOCK
Bug: 258876786
gcip: don't fail dma_fence_init on long name
Bug: 264220687 (repeat)
gcip: implement gcip_dma_fence_show
Bug: 264220687 (repeat)
gcip: implement DMA fence status and signal
Bug: 264220687 (repeat)
gcip: add gcip-dma-fence.c
Bug: 258876786 (repeat)
edgetpu: Add edgetpu_thermal_restore
Bug: 262790767
edgetpu: Add trylock to edgetpu_pm_get_if_powered
Bug: 262790767 (repeat)
edgetpu: set no limit on DMA segment size
Bug: 263431665
edgetpu: Remove unnecessary LINUX_VERSION checks

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

18 files changed, 361 insertions, 104 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index c9fe855..5c20205 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -3,7 +3,7 @@
 # Use the absolute path of this Makefile to get the source directory.
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 GIT_PATH=$(CURRENT_DIR)/../../
-ccflags-y	+= -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -DCONFIG_GCIP=1 \
+ccflags-y	+= $(EXTRA_CFLAGS) -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -DCONFIG_GCIP=1 \
 		   -I$(CURRENT_DIR)/include \
 		   -I$(CURRENT_DIR)/gcip-kernel-driver/include
 
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index c501d7d..8bf8ab5 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -488,6 +488,8 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 		etdev_warn(etdev, "debug dump init fail: %d", ret);
 
 	edgetpu_chip_init(etdev);
+	/* No limit on DMA segment size */
+	dma_set_max_seg_size(etdev->dev, UINT_MAX);
 	return 0;
 
 remove_usage_stats:
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 2b794f1..d6a438d 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -521,13 +521,6 @@ static const struct dma_fence_ops edgetpu_dma_fence_ops = {
 	.release = edgetpu_dma_fence_release,
 };
 
-/* the data type of fence->seqno is u64 in 5.1 */
-#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 1, 0)
-#define SEQ_FMT "%u"
-#else
-#define SEQ_FMT "%llu"
-#endif
-
 int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
 			      struct edgetpu_create_sync_fence_data *datap)
 {
@@ -598,10 +591,8 @@ static int _edgetpu_sync_fence_signal(struct dma_fence *fence, int errno, bool i
 		ret = ignore_signaled ? 0 : -EINVAL;
 		goto out_unlock;
 	}
-	pr_debug("%s: %s-%s%llu-" SEQ_FMT " errno=%d\n", __func__,
-		 fence->ops->get_driver_name(fence),
-		 fence->ops->get_timeline_name(fence), fence->context,
-		 fence->seqno, errno);
+	pr_debug("%s: %s-%s%llu-%llu errno=%d\n", __func__, fence->ops->get_driver_name(fence),
+		 fence->ops->get_timeline_name(fence), fence->context, fence->seqno, errno);
 	if (errno)
 		dma_fence_set_error(fence, errno);
 	ret = dma_fence_signal_locked(fence);
@@ -646,10 +637,10 @@ void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group)
 
 		ret = _edgetpu_sync_fence_signal(fence, -EPIPE, true);
 		if (ret)
-			etdev_warn(group->etdev, "error %d signaling fence %s-%s %llu-" SEQ_FMT,
-				   ret, fence->ops->get_driver_name(fence),
-				   fence->ops->get_timeline_name(fence),
-				   fence->context, fence->seqno);
+			etdev_warn(group->etdev, "error %d signaling fence %s-%s %llu-%llu", ret,
+				   fence->ops->get_driver_name(fence),
+				   fence->ops->get_timeline_name(fence), fence->context,
+				   fence->seqno);
 	}
 }
 
@@ -689,10 +680,8 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 		struct dma_fence *fence = &etfence->fence;
 
 		spin_lock_irq(&etfence->lock);
-		seq_printf(s, "%s-%s %llu-" SEQ_FMT " %s",
-			   fence->ops->get_driver_name(fence),
-			   fence->ops->get_timeline_name(fence),
-			   fence->context, fence->seqno,
+		seq_printf(s, "%s-%s %llu-%llu %s", fence->ops->get_driver_name(fence),
+			   fence->ops->get_timeline_name(fence), fence->context, fence->seqno,
 			   sync_status_str(dma_fence_get_status_locked(fence)));
 
 		if (test_bit(DMA_FENCE_FLAG_TIMESTAMP_BIT, &fence->flags)) {
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index 60017d2..008b92a 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -177,11 +177,9 @@ static int edgetpu_firmware_handshake(struct edgetpu_firmware *et_fw)
 	if (ret)
 		etdev_warn(etdev, "telemetry KCI error: %d", ret);
 
-	if (!IS_ERR(etdev->thermal) && etdev->thermal->cooling_state) {
-		ret = edgetpu_thermal_resume(etdev->dev);
-		if (ret)
-			etdev_warn(etdev, "Thermal resume error: %d", ret);
-	}
+	ret = edgetpu_thermal_restore(etdev);
+	if (ret)
+		etdev_warn(etdev, "thermal restore error: %d", ret);
 
 	/* Set debug dump buffer in FW */
 	edgetpu_get_debug_dump(etdev, 0);
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 869fdec..3bd4b2f 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -938,7 +938,7 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (edgetpu_pm_get_if_powered(group->etdev->pm)) {
+	if (edgetpu_pm_get_if_powered(group->etdev->pm, false)) {
 		mutex_lock(&group->lock);
 		ret = edgetpu_mailbox_external_alloc(group, req);
 		if (ret) {
@@ -975,7 +975,7 @@ static int edgetpu_mailbox_external_disable_free(struct edgetpu_client *client)
 	group = edgetpu_device_group_get(client->group);
 	mutex_unlock(&client->group_lock);
 
-	if (edgetpu_pm_get_if_powered(group->etdev->pm)) {
+	if (edgetpu_pm_get_if_powered(group->etdev->pm, false)) {
 		mutex_lock(&group->lock);
 		edgetpu_mailbox_external_disable_free_locked(group);
 		mutex_unlock(&group->lock);
diff --git a/drivers/edgetpu/edgetpu-mapping.h b/drivers/edgetpu/edgetpu-mapping.h
index 65547bc..91083e3 100644
--- a/drivers/edgetpu/edgetpu-mapping.h
+++ b/drivers/edgetpu/edgetpu-mapping.h
@@ -9,6 +9,7 @@
 
 #include <linux/device.h>
 #include <linux/dma-direction.h>
+#include <linux/dma-map-ops.h>
 #include <linux/iommu.h>
 #include <linux/mutex.h>
 #include <linux/rbtree.h>
@@ -16,9 +17,6 @@
 #include <linux/seq_file.h>
 #include <linux/types.h>
 #include <linux/version.h>
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 10, 0)
-#include <linux/dma-map-ops.h>
-#endif
 
 #include "edgetpu-internal.h"
 #include "edgetpu-mmu.h"
diff --git a/drivers/edgetpu/edgetpu-mmu.h b/drivers/edgetpu/edgetpu-mmu.h
index 44eb970..29bf1a9 100644
--- a/drivers/edgetpu/edgetpu-mmu.h
+++ b/drivers/edgetpu/edgetpu-mmu.h
@@ -16,12 +16,6 @@
 #include "edgetpu-internal.h"
 #include "edgetpu.h"
 
-#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 2, 0)
-#ifndef IOMMU_PASID_INVALID
-#define IOMMU_PASID_INVALID	(-1U)
-#endif
-#endif
-
 /* flags for MMU operations */
 
 /* Whether the TPU address to be allocated can be 64-bit wide. */
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index b6a4542..71cb450 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -71,7 +71,7 @@ void edgetpu_pm_unlock(struct edgetpu_pm *etpm)
 	mutex_unlock(&etpm->p->lock);
 }
 
-bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm)
+bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm, bool trylock)
 {
 	bool ret;
 
@@ -80,12 +80,21 @@ bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm)
 	/* fast fail without holding the lock */
 	if (!etpm->p->power_up_count)
 		return false;
-	mutex_lock(&etpm->p->lock);
+
+	if (trylock) {
+		if (!mutex_trylock(&etpm->p->lock))
+			return false;
+	} else {
+		mutex_lock(&etpm->p->lock);
+	}
+
 	if (etpm->p->power_up_count)
 		ret = !edgetpu_pm_get_locked(etpm);
 	else
 		ret = false;
+
 	mutex_unlock(&etpm->p->lock);
+
 	return ret;
 }
 
diff --git a/drivers/edgetpu/edgetpu-pm.h b/drivers/edgetpu/edgetpu-pm.h
index 1838c75..f0392f0 100644
--- a/drivers/edgetpu/edgetpu-pm.h
+++ b/drivers/edgetpu/edgetpu-pm.h
@@ -59,9 +59,9 @@ void edgetpu_pm_unlock(struct edgetpu_pm *etpm);
  * Caller calls edgetpu_pm_put() to decrease power_up_count if this function
  * returned true, otherwise put() shouldn't be called.
  *
- * Return false if device is not powered, true otherwise.
+ * Return false if device is not powered or trylock fails, true otherwise.
  */
-bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm);
+bool edgetpu_pm_get_if_powered(struct edgetpu_pm *etpm, bool trylock);
 /*
  * Increase power_up_count for active state, power up the device if previous
  * power_up_count was zero.
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index bcfed04..b96d961 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -58,11 +58,21 @@ int edgetpu_thermal_resume(struct device *dev);
 /*
  * Sends the thermal throttling KCI if the device is powered.
  *
- * Returns the return value of KCI if the device is powered, otherwise 0.
+ * Returns the return value of KCI if the device is powered, otherwise -EAGAIN.
  */
 int edgetpu_thermal_kci_if_powered(struct edgetpu_dev *etdev, u32 state);
 
 /*
+ * Sends thermal throttling KCI to restore the last thermal state.
+ *
+ * The caller must guarantee the device stays powered up, typically by calling edgetpu_pm_get() or
+ * by calling this function from the power management functions themselves.
+ *
+ * Returns 0 if no thermal throttling required; otherwise the return value of KCI.
+ */
+int edgetpu_thermal_restore(struct edgetpu_dev *etdev);
+
+/*
  * Holds thermal->lock.
  *
  * Does nothing if the thermal management is not supported.
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index c3424ee..ab68776 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -6,8 +6,14 @@
 CONFIG_GCIP		?= m
 obj-$(CONFIG_GCIP)	+= gcip.o
 
-gcip-objs		:= gcip-alloc-helper.o gcip-domain-pool.o gcip-firmware.o \
-			   gcip-image-config.o gcip-kci.o gcip-mailbox.o gcip-mem-pool.o \
+gcip-objs		:= gcip-alloc-helper.o \
+			   gcip-dma-fence.o \
+			   gcip-domain-pool.o \
+			   gcip-firmware.o \
+			   gcip-image-config.o \
+			   gcip-kci.o \
+			   gcip-mailbox.o \
+			   gcip-mem-pool.o \
 			   gcip-telemetry.o
 
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-dma-fence.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-dma-fence.c
new file mode 100644
index 0000000..4f83670
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-dma-fence.c
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * GCIP support of DMA fences.
+ *
+ * Copyright (C) 2023 Google LLC
+ */
+
+#include <linux/bitops.h>
+#include <linux/device.h>
+#include <linux/dma-fence.h>
+#include <linux/err.h>
+#include <linux/file.h>
+#include <linux/ktime.h>
+#include <linux/list.h>
+#include <linux/seq_file.h>
+#include <linux/spinlock.h>
+#include <linux/string.h>
+#include <linux/sync_file.h>
+#include <linux/time.h>
+
+#include <gcip/gcip-dma-fence.h>
+
+#define to_gfence(fence) container_of(fence, struct gcip_dma_fence, fence)
+
+static int _gcip_dma_fence_signal(struct dma_fence *fence, int error, bool ignore_signaled)
+{
+	int ret;
+
+	if (error > 0)
+		error = -error;
+	if (unlikely(error < -MAX_ERRNO))
+		return -EINVAL;
+
+	spin_lock_irq(fence->lock);
+	/* don't signal fence twice */
+	if (unlikely(test_bit(DMA_FENCE_FLAG_SIGNALED_BIT, &fence->flags))) {
+		ret = ignore_signaled ? 0 : -EBUSY;
+		goto out_unlock;
+	}
+	if (error)
+		dma_fence_set_error(fence, error);
+	ret = dma_fence_signal_locked(fence);
+
+out_unlock:
+	spin_unlock_irq(fence->lock);
+	return ret;
+}
+
+static const char *sync_status_str(int status)
+{
+	if (status < 0)
+		return "error";
+	if (status > 0)
+		return "signaled";
+	return "active";
+}
+
+struct gcip_dma_fence_manager *gcip_dma_fence_manager_create(struct device *dev)
+{
+	struct gcip_dma_fence_manager *mgr = devm_kzalloc(dev, sizeof(*mgr), GFP_KERNEL);
+
+	if (!mgr)
+		return ERR_PTR(-ENOMEM);
+
+	INIT_LIST_HEAD(&mgr->fence_list_head);
+	spin_lock_init(&mgr->fence_list_lock);
+	mgr->dev = dev;
+
+	return mgr;
+}
+
+const char *gcip_dma_fence_get_timeline_name(struct dma_fence *fence)
+{
+	struct gcip_dma_fence *gfence = to_gfence(fence);
+
+	return gfence->timeline_name;
+}
+
+bool gcip_dma_fence_always_true(struct dma_fence *fence)
+{
+	return true;
+}
+
+int gcip_dma_fence_init(struct gcip_dma_fence_manager *mgr, struct gcip_dma_fence *gfence,
+			struct gcip_dma_fence_data *data)
+{
+	unsigned long flags;
+	int fd;
+	struct sync_file *sync_file;
+	int ret;
+
+	strscpy(gfence->timeline_name, data->timeline_name, GCIP_FENCE_TIMELINE_NAME_LEN);
+
+	spin_lock_init(&gfence->lock);
+	INIT_LIST_HEAD(&gfence->fence_list);
+	gfence->mgr = mgr;
+
+	dma_fence_init(&gfence->fence, data->ops, &gfence->lock, dma_fence_context_alloc(1),
+		       data->seqno);
+	GCIP_DMA_FENCE_LIST_LOCK(mgr, flags);
+	list_add_tail(&gfence->fence_list, &mgr->fence_list_head);
+	GCIP_DMA_FENCE_LIST_UNLOCK(mgr, flags);
+
+	if (data->after_init) {
+		ret = data->after_init(gfence);
+		if (ret) {
+			dev_err(mgr->dev, "DMA fence init failed on after_init: %d", ret);
+			goto err_put_fence;
+		}
+	}
+	fd = get_unused_fd_flags(O_CLOEXEC);
+	if (fd < 0) {
+		ret = fd;
+		dev_err(mgr->dev, "Failed to get FD: %d", ret);
+		goto err_put_fence;
+	}
+	sync_file = sync_file_create(&gfence->fence);
+	if (!sync_file) {
+		dev_err(mgr->dev, "Failed to create sync file");
+		ret = -ENOMEM;
+		goto err_put_fd;
+	}
+	/* sync_file holds the reference to fence, so we can drop our reference. */
+	dma_fence_put(&gfence->fence);
+
+	fd_install(fd, sync_file->file);
+	data->fence = fd;
+	return 0;
+
+err_put_fd:
+	put_unused_fd(fd);
+err_put_fence:
+	dma_fence_put(&gfence->fence);
+	return ret;
+}
+
+void gcip_dma_fence_exit(struct gcip_dma_fence *gfence)
+{
+	unsigned long flags;
+
+	GCIP_DMA_FENCE_LIST_LOCK(gfence->mgr, flags);
+	list_del(&gfence->fence_list);
+	GCIP_DMA_FENCE_LIST_UNLOCK(gfence->mgr, flags);
+}
+
+int gcip_dma_fence_status(int fence, int *status)
+{
+	struct dma_fence *fencep;
+
+	fencep = sync_file_get_fence(fence);
+	if (!fencep)
+		return -EBADF;
+	*status = dma_fence_get_status(fencep);
+	dma_fence_put(fencep);
+	return 0;
+}
+
+int gcip_dma_fence_signal(int fence, int error, bool ignore_signaled)
+{
+	struct dma_fence *fencep;
+	int ret;
+
+	fencep = sync_file_get_fence(fence);
+	if (!fencep)
+		return -EBADF;
+	ret = _gcip_dma_fence_signal(fencep, error, ignore_signaled);
+	dma_fence_put(fencep);
+	return ret;
+}
+
+int gcip_dma_fenceptr_signal(struct gcip_dma_fence *gfence, int error, bool ignore_signaled)
+{
+	return _gcip_dma_fence_signal(&gfence->fence, error, ignore_signaled);
+}
+
+void gcip_dma_fence_show(struct gcip_dma_fence *gfence, struct seq_file *s)
+{
+	struct dma_fence *fence = &gfence->fence;
+
+	spin_lock_irq(&gfence->lock);
+
+	seq_printf(s, "%s-%s %llu-%llu %s", fence->ops->get_driver_name(fence),
+		   fence->ops->get_timeline_name(fence), fence->context, fence->seqno,
+		   sync_status_str(dma_fence_get_status_locked(fence)));
+
+	if (test_bit(DMA_FENCE_FLAG_TIMESTAMP_BIT, &fence->flags)) {
+		struct timespec64 ts = ktime_to_timespec64(fence->timestamp);
+
+		seq_printf(s, " @%lld.%09ld", (s64)ts.tv_sec, ts.tv_nsec);
+	}
+
+	if (fence->error)
+		seq_printf(s, " err=%d", fence->error);
+
+	spin_unlock_irq(&gfence->lock);
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
index 312bbdc..62acd0b 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
@@ -12,28 +12,6 @@
 
 #include <gcip/gcip-image-config.h>
 
-#define ADDR_SHIFT 12
-#define SIZE_MODE_BIT BIT(ADDR_SHIFT - 1)
-#define SECURE_SIZE_MASK (SIZE_MODE_BIT - 1u)
-#define NS_SIZE_MASK (BIT(ADDR_SHIFT) - 1u)
-#define ADDR_MASK ~(BIT(ADDR_SHIFT) - 1u)
-
-/* used by ns_iommu_mappings */
-#define CONFIG_TO_MBSIZE(a) (((a) & NS_SIZE_MASK) << 20)
-
-/* used by iommu_mappings */
-static inline __u32 config_to_size(__u32 cfg)
-{
-	__u32 page_size;
-
-	if (cfg & SIZE_MODE_BIT)
-		page_size = cfg & SECURE_SIZE_MASK;
-	else
-		page_size = BIT(cfg & SECURE_SIZE_MASK);
-
-	return page_size << PAGE_SHIFT;
-}
-
 static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 				struct gcip_image_config *config)
 {
@@ -49,8 +27,8 @@ static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 			ret = -EIO;
 			goto err;
 		}
-		size = config_to_size(config->iommu_mappings[i].image_config_value);
-		paddr = config->iommu_mappings[i].image_config_value & ADDR_MASK;
+		size = gcip_config_to_size(config->iommu_mappings[i].image_config_value);
+		paddr = config->iommu_mappings[i].image_config_value & GCIP_IMG_CFG_ADDR_MASK;
 
 		dev_dbg(parser->dev, "Image config adding IOMMU mapping: %pad -> %pap", &daddr,
 			&paddr);
@@ -74,7 +52,7 @@ static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 err:
 	while (i--) {
 		daddr = config->iommu_mappings[i].virt_address;
-		size = config_to_size(config->iommu_mappings[i].image_config_value);
+		size = gcip_config_to_size(config->iommu_mappings[i].image_config_value);
 		parser->ops->unmap(parser->data, daddr, size, GCIP_IMAGE_CONFIG_FLAGS_SECURE);
 	}
 	return ret;
@@ -89,7 +67,7 @@ static void clear_iommu_mappings(struct gcip_image_config_parser *parser,
 
 	for (i = config->num_iommu_mappings - 1; i >= 0; i--) {
 		daddr = config->iommu_mappings[i].virt_address;
-		size = config_to_size(config->iommu_mappings[i].image_config_value);
+		size = gcip_config_to_size(config->iommu_mappings[i].image_config_value);
 		dev_dbg(parser->dev, "Image config removing IOMMU mapping: %pad size=%#lx", &daddr,
 			size);
 		parser->ops->unmap(parser->data, daddr, size, GCIP_IMAGE_CONFIG_FLAGS_SECURE);
@@ -105,13 +83,13 @@ static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 	phys_addr_t paddr = 0;
 
 	for (i = 0; i < config->num_ns_iommu_mappings; i++) {
-		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
+		daddr = config->ns_iommu_mappings[i] & GCIP_IMG_CFG_ADDR_MASK;
 		if (unlikely(!daddr)) {
 			dev_warn(parser->dev, "Invalid config, device address is zero");
 			ret = -EIO;
 			goto err;
 		}
-		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
+		size = gcip_ns_config_to_size(config->ns_iommu_mappings[i]);
 		dev_dbg(parser->dev, "Image config adding NS IOMMU mapping: %pad -> %pap", &daddr,
 			&paddr);
 		if (unlikely(daddr + size <= daddr || paddr + size <= paddr)) {
@@ -128,8 +106,8 @@ static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 
 err:
 	while (i--) {
-		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
-		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
+		size = gcip_ns_config_to_size(config->ns_iommu_mappings[i]);
+		daddr = config->ns_iommu_mappings[i] & GCIP_IMG_CFG_ADDR_MASK;
 		parser->ops->unmap(parser->data, daddr, size, 0);
 	}
 	return ret;
@@ -143,8 +121,8 @@ static void clear_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 	int i;
 
 	for (i = config->num_ns_iommu_mappings - 1; i >= 0; i--) {
-		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
-		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
+		size = gcip_ns_config_to_size(config->ns_iommu_mappings[i]);
+		daddr = config->ns_iommu_mappings[i] & GCIP_IMG_CFG_ADDR_MASK;
 		dev_dbg(parser->dev, "Image config removing NS IOMMU mapping: %pad size=%#lx",
 			&daddr, size);
 		parser->ops->unmap(parser->data, daddr, size, 0);
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-dma-fence.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-dma-fence.h
index a46bcbb..c0a7d68 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-dma-fence.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-dma-fence.h
@@ -18,7 +18,7 @@
 
 /* Used before accessing the list headed by mgr->fence_list_head. */
 #define GCIP_DMA_FENCE_LIST_LOCK(mgr, flags) spin_lock_irqsave(&mgr->fence_list_lock, flags)
-#define GCIP_DMA_FENCE_LIST_UNLOCK(mgr, flags) spin_lock_irqrestore(&mgr->fence_list_lock, flags)
+#define GCIP_DMA_FENCE_LIST_UNLOCK(mgr, flags) spin_unlock_irqrestore(&mgr->fence_list_lock, flags)
 
 /*
  * A macro to loop through all fences under a gcip_dma_fence_manager.
@@ -30,6 +30,8 @@
 #define gcip_for_each_fence(mgr, gfence)                                                           \
 	list_for_each_entry(gfence, &mgr->fence_list_head, fence_list)
 
+#define to_gcip_fence(fence) container_of(fence, struct gcip_dma_fence, fence)
+
 struct gcip_dma_fence_manager {
 	/* The list of all fence objects for debugging. */
 	struct list_head fence_list_head;
@@ -57,7 +59,9 @@ struct gcip_dma_fence_data {
 	 * the gcip_dma_fence_init() call.
 	 */
 	char *timeline_name;
-	/* The DMA fence operators to initialize the fence with. */
+	/*
+	 * The DMA fence operators to initialize the fence with.
+	 */
 	const struct dma_fence_ops *ops;
 	/* The sequence number to initialize the fence with. */
 	u32 seqno;
@@ -95,6 +99,20 @@ bool gcip_dma_fence_always_true(struct dma_fence *fence);
 
 /* End of helpers for setting dma_fence_ops. */
 
+/*
+ * This function does
+ *  1. Initialize the DMA fence object
+ *  2. Call after_init() if present
+ *  3. Install an FD associates to the created DMA fence
+ *
+ * This function never fails on step 1, so this function returns an error only if after_init() fails
+ * (step 2) or FD allocation fails (step 3).
+ * In either failure case, @ops->release is always called. Therefore @ops->release may need to
+ * distinguish whether after_init() succeeded.
+ *
+ * It's always safe to call gcip_dma_fence_exit() in @ops->release because that function reverts
+ * step 1.
+ */
 int gcip_dma_fence_init(struct gcip_dma_fence_manager *mgr, struct gcip_dma_fence *gfence,
 			struct gcip_dma_fence_data *data);
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
index a995188..1d00ef9 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
@@ -8,6 +8,8 @@
 #ifndef __GCIP_IMAGE_CONFIG_H__
 #define __GCIP_IMAGE_CONFIG_H__
 
+#include <asm/page.h>
+#include <linux/bits.h>
 #include <linux/types.h>
 
 #define GCIP_FW_NUM_VERSIONS 4
@@ -80,6 +82,39 @@ struct gcip_image_config_parser {
 	struct gcip_image_config last_config;
 };
 
+#define GCIP_IMG_CFG_ADDR_SHIFT 12
+#define GCIP_IMG_CFG_MB_SHIFT 20
+#define GCIP_IMG_CFG_SIZE_MODE_BIT BIT(GCIP_IMG_CFG_ADDR_SHIFT - 1)
+#define GCIP_IMG_CFG_SECURE_SIZE_MASK (GCIP_IMG_CFG_SIZE_MODE_BIT - 1u)
+#define GCIP_IMG_CFG_NS_SIZE_MASK (GCIP_IMG_CFG_SIZE_MODE_BIT - 1u)
+#define GCIP_IMG_CFG_ADDR_MASK ~(BIT(GCIP_IMG_CFG_ADDR_SHIFT) - 1u)
+
+/* For decoding the size of ns_iommu_mappings. */
+static inline u32 gcip_ns_config_to_size(u32 cfg)
+{
+	u32 size;
+
+	if (cfg & GCIP_IMG_CFG_SIZE_MODE_BIT)
+		size = (cfg & GCIP_IMG_CFG_NS_SIZE_MASK) << PAGE_SHIFT;
+	else
+		size = (cfg & GCIP_IMG_CFG_NS_SIZE_MASK) << GCIP_IMG_CFG_MB_SHIFT;
+
+	return size;
+}
+
+/* For decoding the size of iommu_mappings. */
+static inline u32 gcip_config_to_size(u32 cfg)
+{
+	u32 page_size;
+
+	if (cfg & GCIP_IMG_CFG_SIZE_MODE_BIT)
+		page_size = cfg & GCIP_IMG_CFG_SECURE_SIZE_MASK;
+	else
+		page_size = BIT(cfg & GCIP_IMG_CFG_SECURE_SIZE_MASK);
+
+	return page_size << PAGE_SHIFT;
+}
+
 /*
  * Initializes the image configuration parser.
  *
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 bda1b40..3649779 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-kci.h
@@ -109,6 +109,11 @@ enum gcip_kci_code {
  */
 enum gcip_reverse_kci_code {
 	GCIP_RKCI_CHIP_CODE_FIRST = 0,
+	GCIP_RKCI_PM_QOS_REQUEST = GCIP_RKCI_CHIP_CODE_FIRST + 1,
+	GCIP_RKCI_CHANGE_BTS_SCENARIO = GCIP_RKCI_CHIP_CODE_FIRST + 2,
+	GCIP_RKCI_PM_QOS_BTS_REQUEST = GCIP_RKCI_CHIP_CODE_FIRST + 3,
+	GCIP_RKCI_DSP_CORE_TELEMETRY_TRY_READ = GCIP_RKCI_CHIP_CODE_FIRST + 4,
+	GCIP_RKCI_CLIENT_FATAL_ERROR_NOTIFY = GCIP_RKCI_CHIP_CODE_FIRST + 5,
 	GCIP_RKCI_CHIP_CODE_LAST = 0x7FFF,
 	GCIP_RKCI_GENERIC_CODE_FIRST = 0x8000,
 	GCIP_RKCI_FIRMWARE_CRASH = GCIP_RKCI_GENERIC_CODE_FIRST + 0,
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index 8fb5288..e4f1aaf 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -314,7 +314,7 @@ long edgetpu_soc_pm_get_rate(struct edgetpu_dev *etdev, int flags)
 		return -EINVAL;
 
 	/* We need to keep TPU being powered to ensure CMU read is valid. */
-	if (!edgetpu_pm_get_if_powered(etdev->pm))
+	if (!edgetpu_pm_get_if_powered(etdev->pm, false))
 		return 0;
 	pll_con3 = readl(cmu_base + PLL_CON3_OFFSET);
 	edgetpu_pm_put(etdev->pm);
diff --git a/drivers/edgetpu/mobile-thermal.c b/drivers/edgetpu/mobile-thermal.c
index df3a6e0..4794053 100644
--- a/drivers/edgetpu/mobile-thermal.c
+++ b/drivers/edgetpu/mobile-thermal.c
@@ -130,11 +130,7 @@ static int edgetpu_state2power_internal(unsigned long state, u32 *power,
 	return -EINVAL;
 }
 
-static int edgetpu_get_requested_power(struct thermal_cooling_device *cdev,
-#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 10, 0)
-				       struct thermal_zone_device *tz,
-#endif
-				       u32 *power)
+static int edgetpu_get_requested_power(struct thermal_cooling_device *cdev, u32 *power)
 {
 	unsigned long state_original;
 	struct edgetpu_thermal *thermal = cdev->devdata;
@@ -143,11 +139,7 @@ static int edgetpu_get_requested_power(struct thermal_cooling_device *cdev,
 	return edgetpu_state2power_internal(state_original, power, thermal);
 }
 
-static int edgetpu_state2power(struct thermal_cooling_device *cdev,
-#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 10, 0)
-			       struct thermal_zone_device *tz,
-#endif
-			       unsigned long state, u32 *power)
+static int edgetpu_state2power(struct thermal_cooling_device *cdev, unsigned long state, u32 *power)
 {
 	struct edgetpu_thermal *thermal = cdev->devdata;
 
@@ -159,11 +151,7 @@ static int edgetpu_state2power(struct thermal_cooling_device *cdev,
 	return edgetpu_state2power_internal(state_pwr_map[state].state, power, thermal);
 }
 
-static int edgetpu_power2state(struct thermal_cooling_device *cdev,
-#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 10, 0)
-			       struct thermal_zone_device *tz,
-#endif
-			       u32 power, unsigned long *state)
+static int edgetpu_power2state(struct thermal_cooling_device *cdev, u32 power, unsigned long *state)
 {
 	int i, penultimate_throttle_state;
 	struct edgetpu_thermal *thermal = cdev->devdata;
@@ -418,21 +406,52 @@ int edgetpu_thermal_resume(struct device *dev)
 	return ret;
 }
 
+static int edgetpu_thermal_kci(struct edgetpu_dev *etdev, u32 state)
+{
+	int ret;
+
+	edgetpu_kci_block_bus_speed_control(etdev, true);
+
+	ret = edgetpu_kci_notify_throttling(etdev, state);
+	if (ret)
+		etdev_err_ratelimited(etdev, "Failed to notify FW about power state %u, error:%d",
+				      state, ret);
+
+	edgetpu_kci_block_bus_speed_control(etdev, false);
+
+	return ret;
+}
+
 int edgetpu_thermal_kci_if_powered(struct edgetpu_dev *etdev, u32 state)
 {
-	int ret = -EAGAIN;
+	int ret;
 
-	if (edgetpu_pm_get_if_powered(etdev->pm)) {
-		edgetpu_kci_block_bus_speed_control(etdev, true);
+	/* Use trylock since this function might be called during power up. */
+	if (!edgetpu_pm_get_if_powered(etdev->pm, true))
+		return -EAGAIN;
 
-		ret = edgetpu_kci_notify_throttling(etdev, state);
-		if (ret)
-			etdev_err_ratelimited(etdev,
-					      "Failed to notify FW about power state %u, error:%d",
-					      state, ret);
+	ret = edgetpu_thermal_kci(etdev, state);
+	edgetpu_pm_put(etdev->pm);
+
+	return ret;
+}
+
+int edgetpu_thermal_restore(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_thermal *thermal = etdev->thermal;
+	int ret = 0;
+
+	if (IS_ERR_OR_NULL(thermal))
+		return 0;
+
+	edgetpu_thermal_lock(thermal);
+
+	if (edgetpu_thermal_is_suspended(thermal))
+		ret = edgetpu_thermal_kci(etdev, TPU_OFF);
+	else if (thermal->cooling_state)
+		ret = edgetpu_thermal_kci(etdev, state_pwr_map[thermal->cooling_state].state);
+
+	edgetpu_thermal_unlock(thermal);
 
-		edgetpu_kci_block_bus_speed_control(etdev, false);
-		edgetpu_pm_put(etdev->pm);
-	}
 	return ret;
 }