[Copybara Auto Merge] Merge branch whitechapel into partner-androidandroid-u-beta-1_r0.5android-u-beta-1_r0.4android-u-beta-1_r0.3android-gs-raviole-5.10-u-beta1android-gs-pantah-5.10-u-beta1android-gs-bluejay-5.10-u-beta1

edgetpu: whitechapel: sync with darwinn-2.0
bedb92a1c edgetpu: Add missing pm error handling
fbbd28615 edgetpu: set no limit on DMA segment size
305c1ccf2 edgetpu: Fix missing wakelock unlock
Bug: 263431665
edgetpu: sync with darwinn-2.0
df80d447c edgetpu: signal any unsignaled dma fences at device group disband
41731ed12 edgetpu: dma fence debugfs show owning group
765a088dd edgetpu: dma fence creation optionally associated with device group
4c3abd990 edgetpu: Enable best-fit algo for iova allocator
93c1f3241 edgetpu: Add a static debug pointer to driver state
3b88f37a6 edgetpu: Add reasonable defaults to Makefile
7bc202cd6 edgetpu: Fix pmu-status-base reading
e00d0a155 edgetpu: Return no space error from dma_map_sg
f6fecb36e edgetpu: introduce PMU STATUS query
628ffc1bb edgetpu: use rate limiting for reverse KCI failure warning
87add7e7a edgetpu: janeiro: Check block status before power-up
f02bb3f28 edgetpu: warn if KCI get usage fails
Bug: 258868303
Bug: 243837563
Bug: 260768055
Bug: 249727851
Bug: 243837563 (repeat)
Bug: 193072630
edgetpu: abrolhos: Fix include file name to gs_tmu.h
commit 9adbf68cf96 ("drivers: thermal: Rename gs101_tmu to gs_tmu")
renamed including file from gs101_tmu.h to gs_tmu.h. Rename including
path to avoid build breakage.
Bug: 238399460
edgetpu: abrolhos: add pre-allocated iommu domains
Add pre-allocated iommu domains per contexts.
Bug: 251639928
(cherry picked from commit 3778a777b969d38a2e2443c64d84e36bbe14e47f)
edgetpu: fix invalid paging on invalid map size
Bug: 245095356
edgetpu: rename mobile PM functions to be edgetpu-specific
edgetpu: check mobile firmware header magic, warn if mismatch
edgetpu: add error message on partial buffer map
edgetpu: remove KCI unmap buffer code
edgetpu: dma-fence debugfs show use dma_fence_ops
edgetpu: remove unnecessary overflow check in buffer map code
edgetpu: use mm_sem instead of mmap_lock on older Linux
Bug: 238359788
edgetpu: remove include of mmap_lock.h
Bug: 238359788 (repeat)
edgetpu: hold mmap lock around call to find_extend_vma
Bug: 237404338
edgetpu: Remove next-gen chipset from darwinn-2.0
Bug: 236213851
edgetpu: hermosa: set KCI timeout to 5 seconds
Bug: 237004877
edgetpu: make KCI source optionally customizable
Bug: 237004877 (repeat)
edgetpu: unittests: Fix __alloc_pages_nodemask warning
edgetpu: reduce KCI wait timeout.
Bug: 234458631
edgetpu: gitignore: ignore gcip-kernel-driver
Bug: 234255813
edgetpu: remove warning about 64-bit addressing on mobile
edgetpu: mobile: deprecate the code for separate pmqos and bts request.
Bug: 233338364
edgetpu: respond to reverse KCI commands.
Bug: 229143599
edgetpu: add debugfs "wakelock" power up/down control
Bug: 232449109
GitOrigin-RevId: df8bf721bc0cda8429a23d66e8d912a32f2848f3
Change-Id: Ib1a7a9eb29f5f6d0cd240ed075b2a8eaabd7d3f5

21 files changed, 307 insertions, 148 deletions

diff --git a/drivers/edgetpu/.gitignore b/drivers/edgetpu/.gitignore
new file mode 100644
index 0000000..c04c6ad
--- /dev/null
+++ b/drivers/edgetpu/.gitignore
@@ -0,0 +1 @@
+/gcip-kernel-driver
diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index e0c3588..fab3cd2 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -10,10 +10,11 @@ else
 	ccflags-y	+= -DGIT_REPO_TAG=\"Not\ a\ git\ repository\"
 endif
 
-edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-domain-pool.o
+edgetpu-objs	:= edgetpu-mailbox.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-domain-pool.o
 
-abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o abrolhos-iommu.o abrolhos-debug-dump.o abrolhos-usage-stats.o abrolhos-wakelock.o $(edgetpu-objs)
+mobile-objs	:= edgetpu-kci.o
 
+abrolhos-y	:= abrolhos-device.o abrolhos-device-group.o abrolhos-fs.o abrolhos-core.o abrolhos-platform.o abrolhos-firmware.o abrolhos-thermal.o abrolhos-pm.o abrolhos-iommu.o abrolhos-debug-dump.o abrolhos-usage-stats.o abrolhos-wakelock.o $(mobile-objs) $(edgetpu-objs)
 
 
 
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index ac850d2..fbdba93 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -14,21 +14,25 @@ else
 endif
 
 edgetpu-objs	:= edgetpu-async.o edgetpu-dmabuf.o edgetpu-iremap-pool.o \
-		   edgetpu-kci.o edgetpu-mailbox.o edgetpu-mapping.o \
+		   edgetpu-mailbox.o edgetpu-mapping.o \
 		   edgetpu-sw-watchdog.o edgetpu-telemetry.o \
 		   edgetpu-firmware-util.o edgetpu-firmware.o \
 		   edgetpu-domain-pool.o
 
+mobile-objs	:= edgetpu-kci.o
+
 abrolhos-objs	:= abrolhos-core.o abrolhos-debug-dump.o \
 		   abrolhos-device-group.o abrolhos-device.o \
 		   abrolhos-firmware.o abrolhos-fs.o abrolhos-iommu.o \
 		   abrolhos-platform.o abrolhos-pm.o abrolhos-thermal.o \
 		   abrolhos-usage-stats.o abrolhos-wakelock.o \
-		   $(edgetpu-objs)
-
+		   $(mobile-objs) $(edgetpu-objs)
 
 
 KBUILD_OPTIONS += CONFIG_ABROLHOS=m
 
+KERNEL_SRC ?= /lib/modules/$(shell uname -r)/build
+M ?= $(shell pwd)
+
 modules modules_install clean:
 	$(MAKE) -C $(KERNEL_SRC) M=$(M) W=1 $(KBUILD_OPTIONS) $(@)
diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index 0aecdd1..bc1108f 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -93,18 +93,12 @@ void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
 				     struct edgetpu_kci_response_element *resp)
 {
 	switch (resp->code) {
-	case RKCI_CODE_PM_QOS:
-		mobile_pm_set_pm_qos(etdev, resp->retval);
-		break;
-	case RKCI_CODE_BTS:
-		mobile_pm_set_bts(etdev, resp->retval);
-		break;
 	case RKCI_CODE_PM_QOS_BTS:
 		/* FW indicates to ignore the request by setting them to undefined values. */
 		if (resp->retval != (typeof(resp->retval))~0ull)
-			mobile_pm_set_pm_qos(etdev, resp->retval);
+			edgetpu_mobile_pm_set_pm_qos(etdev, resp->retval);
 		if (resp->status != (typeof(resp->status))~0ull)
-			mobile_pm_set_bts(etdev, resp->status);
+			edgetpu_mobile_pm_set_bts(etdev, resp->status);
 		break;
 	default:
 		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n",
diff --git a/drivers/edgetpu/abrolhos-pm.c b/drivers/edgetpu/abrolhos-pm.c
index 8bf40e2..015e755 100644
--- a/drivers/edgetpu/abrolhos-pm.c
+++ b/drivers/edgetpu/abrolhos-pm.c
@@ -85,5 +85,5 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 	platform_pwr->after_create = abrolhos_pm_after_create;
 	platform_pwr->acpm_set_rate = exynos_acpm_set_rate;
 
-	return mobile_pm_create(etdev);
+	return edgetpu_mobile_pm_create(etdev);
 }
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index ff1ad77..8d66b81 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -24,4 +24,10 @@
 #define EDGETPU_NUM_CORES 1
 #endif
 
+/* Uses a smaller size for unittests to avoid DMA warnings. */
+#if IS_ENABLED(CONFIG_EDGETPU_TEST)
+#undef EDGETPU_DEBUG_DUMP_MEM_SIZE
+#define EDGETPU_DEBUG_DUMP_MEM_SIZE 0x20000
+#endif
+
 #endif /* __EDGETPU_CONFIG_H__ */
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 81c54a6..6973bdb 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -490,6 +490,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_kci:
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 86ef9d2..7734f33 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -27,6 +27,7 @@
 #include "edgetpu-async.h"
 #include "edgetpu-config.h"
 #include "edgetpu-device-group.h"
+#include "edgetpu-dmabuf.h"
 #include "edgetpu-dram.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -505,6 +506,8 @@ static void edgetpu_device_group_release(struct edgetpu_device_group *group)
 		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
 		edgetpu_mmu_free_domain(group->etdev, group->etdomain);
 	}
+	/* Signal any unsignaled dma fences owned by the group with an error. */
+	edgetpu_sync_fence_group_shutdown(group);
 	group->status = EDGETPU_DEVICE_GROUP_DISBANDED;
 }
 
@@ -721,6 +724,7 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 	group->vii.etdev = client->etdev;
 	mutex_init(&group->lock);
 	rwlock_init(&group->events.lock);
+	INIT_LIST_HEAD(&group->dma_fence_list);
 	edgetpu_mapping_init(&group->host_mappings);
 	edgetpu_mapping_init(&group->dmabuf_mappings);
 	group->mbox_attr = *attr;
@@ -1174,12 +1178,9 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 		return ERR_PTR(-EFAULT);
 	}
 	offset = host_addr & (PAGE_SIZE - 1);
-	/* overflow check (should also be caught by access_ok) */
-	if (unlikely((size + offset) / PAGE_SIZE >= UINT_MAX - 1 || size + offset < size)) {
-		etdev_err(etdev, "address overflow in buffer map request");
-		return ERR_PTR(-EFAULT);
-	}
 	num_pages = DIV_ROUND_UP((size + offset), PAGE_SIZE);
+	if (num_pages * PAGE_SIZE < size + offset)
+		return ERR_PTR(-EINVAL);
 	etdev_dbg(etdev, "%s: hostaddr=%#llx pages=%u", __func__, host_addr, num_pages);
 	/*
 	 * "num_pages" is decided from user-space arguments, don't show warnings
@@ -1196,6 +1197,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	 * it with FOLL_WRITE.
 	 * default to read/write if find_extend_vma returns NULL
 	 */
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	down_read(&current->mm->mmap_sem);
+#else
+	mmap_read_lock(current->mm);
+#endif
 	vma = find_extend_vma(current->mm, host_addr & PAGE_MASK);
 	if (vma && !(vma->vm_flags & VM_WRITE)) {
 		foll_flags &= ~FOLL_WRITE;
@@ -1203,6 +1209,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	} else {
 		*preadonly = false;
 	}
+#if LINUX_VERSION_CODE < KERNEL_VERSION(5, 8, 0)
+	up_read(&current->mm->mmap_sem);
+#else
+	mmap_read_unlock(current->mm);
+#endif
 
 	/* Try fast call first, in case it's actually faster. */
 	ret = pin_user_pages_fast(host_addr & PAGE_MASK, num_pages, foll_flags,
@@ -1265,6 +1276,8 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 			  "pin_user_pages partial %u:%pK npages=%u pinned=%d",
 			  group->workload_id, (void *)host_addr, num_pages,
 			  ret);
+		etdev_err(etdev, "can only lock %u of %u pages requested",
+			  (unsigned int)ret, num_pages);
 		num_pages = ret;
 		ret = -EFAULT;
 		goto error;
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index bf2776b..fffcf75 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -129,6 +129,9 @@ struct edgetpu_device_group {
 	/* Mask of errors set for this group. */
 	uint fatal_errors;
 
+	/* List of DMA fences owned by this group */
+	struct list_head dma_fence_list;
+
 	/* end of fields protected by @lock */
 
 	/* TPU IOVA mapped to host DRAM space */
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 657ae75..72072f6 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -72,6 +72,9 @@ struct edgetpu_dmabuf_map {
  * @fence:		the base DMA fence
  * @lock:		spinlock protecting updates to @fence
  * @timeline_name:	name of the timeline associated with the fence
+ * @group:		owning device group
+ * @etfence_list:	global list of all edgetpu DMA fences
+ * @group_list:		list of DMA fences owned by the same group
  *
  * It is likely timelines will become a separate object in the future,
  * but for now there's a unique named timeline associated with each fence.
@@ -80,7 +83,9 @@ struct edgetpu_dma_fence {
 	struct dma_fence fence;
 	spinlock_t lock;
 	char timeline_name[EDGETPU_SYNC_TIMELINE_NAME_LEN];
+	struct edgetpu_device_group *group;
 	struct list_head etfence_list;
+	struct list_head group_list;
 };
 
 /* List of all edgetpu fence objects for debugging. */
@@ -873,6 +878,7 @@ static const char *edgetpu_dma_fence_get_timeline_name(struct dma_fence *fence)
 static void edgetpu_dma_fence_release(struct dma_fence *fence)
 {
 	struct edgetpu_dma_fence *etfence = to_etfence(fence);
+	struct edgetpu_device_group *group;
 	unsigned long flags;
 
 	if (!etfence)
@@ -881,6 +887,17 @@ static void edgetpu_dma_fence_release(struct dma_fence *fence)
 	spin_lock_irqsave(&etfence_list_lock, flags);
 	list_del(&etfence->etfence_list);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
+
+	/* TODO(b/258868303): Don't remove this check when group required, might not yet be set. */
+	group = etfence->group;
+	if (group) {
+		mutex_lock(&group->lock);
+		list_del(&etfence->group_list);
+		mutex_unlock(&group->lock);
+		/* Release this fence's reference on the owning group. */
+		edgetpu_device_group_put(group);
+	}
+
 	kfree(etfence);
 }
 
@@ -904,7 +921,8 @@ static const struct dma_fence_ops edgetpu_dma_fence_ops = {
 #define SEQ_FMT "%llu"
 #endif
 
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap)
 {
 	int fd = get_unused_fd_flags(O_CLOEXEC);
 	int ret;
@@ -926,6 +944,7 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	 * list_head is needed for list_del().
 	 */
 	INIT_LIST_HEAD(&etfence->etfence_list);
+	INIT_LIST_HEAD(&etfence->group_list);
 	memcpy(&etfence->timeline_name, &datap->timeline_name,
 	       EDGETPU_SYNC_TIMELINE_NAME_LEN - 1);
 
@@ -945,6 +964,14 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	list_add_tail(&etfence->etfence_list, &etfence_list_head);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
 
+	/* TODO(b/258868303): Make group required, disallow creating fence we can't track. */
+	if (group) {
+		etfence->group = edgetpu_device_group_get(group);
+		mutex_lock(&group->lock);
+		list_add_tail(&etfence->group_list, &group->dma_fence_list);
+		mutex_unlock(&group->lock);
+	}
+
 	fd_install(fd, sync_file->file);
 	datap->fence = fd;
 	return 0;
@@ -954,26 +981,14 @@ err_put_fd:
 	return ret;
 }
 
-int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+static int _edgetpu_sync_fence_signal(struct dma_fence *fence, int errno, bool ignore_signaled)
 {
-	struct dma_fence *fence;
-	int errno;
 	int ret;
 
-	errno = datap->error;
-	if (errno > 0)
-		errno = -errno;
-	if (errno < -MAX_ERRNO)
-		return -EINVAL;
-
-	fence = sync_file_get_fence(datap->fence);
-	if (!fence)
-		return -EINVAL;
-
 	spin_lock_irq(fence->lock);
 	/* don't signal fence twice */
 	if (unlikely(test_bit(DMA_FENCE_FLAG_SIGNALED_BIT, &fence->flags))) {
-		ret = -EINVAL;
+		ret = ignore_signaled ? 0 : -EINVAL;
 		goto out_unlock;
 	}
 	pr_debug("%s: %s-%s%llu-" SEQ_FMT " errno=%d\n", __func__,
@@ -986,10 +1001,51 @@ int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
 
 out_unlock:
 	spin_unlock_irq(fence->lock);
+	return ret;
+}
+
+int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+{
+	struct dma_fence *fence;
+	int errno;
+	int ret;
+
+	errno = datap->error;
+	if (errno > 0)
+		errno = -errno;
+	if (errno < -MAX_ERRNO)
+		return -EINVAL;
+
+	fence = sync_file_get_fence(datap->fence);
+	if (!fence)
+		return -EINVAL;
+
+	ret = _edgetpu_sync_fence_signal(fence, errno, false);
 	dma_fence_put(fence);
 	return ret;
 }
 
+/* Caller holds group lock. */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group)
+{
+	struct list_head *pos;
+	int ret;
+
+	lockdep_assert_held(&group->lock);
+	list_for_each(pos, &group->dma_fence_list) {
+		struct edgetpu_dma_fence *etfence =
+			container_of(pos, struct edgetpu_dma_fence, group_list);
+		struct dma_fence *fence = &etfence->fence;
+
+		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);
+	}
+}
+
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap)
 {
 	struct dma_fence *fence;
@@ -1027,8 +1083,9 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 
 		spin_lock_irq(&etfence->lock);
 		seq_printf(s, "%s-%s %llu-" SEQ_FMT " %s",
-			   edgetpu_dma_fence_get_driver_name(fence),
-			   etfence->timeline_name, fence->context, fence->seqno,
+			   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)) {
@@ -1041,7 +1098,9 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 
 		if (fence->error)
 			seq_printf(s, " err=%d", fence->error);
-
+		/* TODO(b/258868303): Remove check when group is required. */
+		if (etfence->group)
+			seq_printf(s, " group=%u", etfence->group->workload_id);
 		seq_putc(s, '\n');
 		spin_unlock_irq(&etfence->lock);
 	}
diff --git a/drivers/edgetpu/edgetpu-dmabuf.h b/drivers/edgetpu/edgetpu-dmabuf.h
index f35b2d3..2b9e281 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.h
+++ b/drivers/edgetpu/edgetpu-dmabuf.h
@@ -38,11 +38,17 @@ int edgetpu_map_bulk_dmabuf(struct edgetpu_device_group *group,
 int edgetpu_unmap_bulk_dmabuf(struct edgetpu_device_group *group,
 			      tpu_addr_t tpu_addr);
 /* Create a DMA sync fence via ioctl */
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap);
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap);
 /* Signal a DMA sync fence, optionally specifying error status */
 int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap);
 /* Return DMA sync fence status */
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap);
+/*
+ * Send error signal to any remaining unsignalled DMA sync fences in a group being disbanded.
+ * Caller holds group lock.
+ */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group);
 /* Dump sync fence info from debugfs */
 int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused);
 
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index d0efb67..68ff32e 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -432,6 +432,7 @@ edgetpu_ioctl_unmap_dmabuf(struct edgetpu_client *client,
 }
 
 static int edgetpu_ioctl_sync_fence_create(
+	struct edgetpu_client *client,
 	struct edgetpu_create_sync_fence_data __user *datap)
 {
 	struct edgetpu_create_sync_fence_data data;
@@ -439,7 +440,13 @@ static int edgetpu_ioctl_sync_fence_create(
 
 	if (copy_from_user(&data, (void __user *)datap, sizeof(data)))
 		return -EFAULT;
-	ret = edgetpu_sync_fence_create(&data);
+	LOCK(client);
+	if (!client->group)
+		/* TODO(b/258868303): Require a group, disallow creating a fence we can't track. */
+		etdev_warn(client->etdev,
+			   "client creating sync fence not joined to a device group");
+	ret = edgetpu_sync_fence_create(client->group, &data);
+	UNLOCK(client);
 	if (ret)
 		return ret;
 	if (copy_to_user((void __user *)datap, &data, sizeof(data)))
@@ -597,48 +604,48 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 		etdev_warn_ratelimited(client->etdev,
 				       "wakelock acquire rejected due to thermal suspend");
 		edgetpu_thermal_unlock(thermal);
-		goto error_unlock;
+		goto error_client_unlock;
 	} else {
 		ret = edgetpu_pm_get(client->etdev->pm);
 		edgetpu_thermal_unlock(thermal);
 		if (ret) {
 			etdev_warn(client->etdev, "%s: pm_get failed (%d)",
 				   __func__, ret);
-			goto error_unlock;
+			goto error_client_unlock;
 		}
 	}
 	edgetpu_wakelock_lock(client->wakelock);
 	/* when NO_WAKELOCK: count should be 1 so here is a no-op */
 	count = edgetpu_wakelock_acquire(client->wakelock);
 	if (count < 0) {
-		edgetpu_pm_put(client->etdev->pm);
 		ret = count;
-		goto error_unlock;
+		goto error_wakelock_unlock;
 	}
 	if (!count) {
 		if (client->group)
 			ret = edgetpu_group_attach_and_open_mailbox(client->group);
 		if (ret) {
-			etdev_warn(client->etdev,
-				   "failed to attach mailbox: %d", ret);
-			edgetpu_pm_put(client->etdev->pm);
+			etdev_warn(client->etdev, "failed to attach mailbox: %d", ret);
 			edgetpu_wakelock_release(client->wakelock);
-			edgetpu_wakelock_unlock(client->wakelock);
-			goto error_unlock;
+			goto error_wakelock_unlock;
 		}
-	} else {
-		/* Balance the power up count due to pm_get above.*/
-		edgetpu_pm_put(client->etdev->pm);
 	}
+
+error_wakelock_unlock:
 	edgetpu_wakelock_unlock(client->wakelock);
+
+	/* Balance the power up count due to pm_get above.*/
+	if (ret || count)
+		edgetpu_pm_put(client->etdev->pm);
+
+error_client_unlock:
 	UNLOCK(client);
-	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
-		  count + 1);
-	return 0;
-error_unlock:
-	UNLOCK(client);
-	etdev_err(client->etdev, "client pid %d failed to acquire wakelock",
-		  client->pid);
+
+	if (ret)
+		etdev_err(client->etdev, "client pid %d failed to acquire wakelock", client->pid);
+	else
+		etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__, count + 1);
+
 	return ret;
 }
 
@@ -786,7 +793,7 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 		ret = edgetpu_ioctl_allocate_device_buffer(client, (u64)argp);
 		break;
 	case EDGETPU_CREATE_SYNC_FENCE:
-		ret = edgetpu_ioctl_sync_fence_create(argp);
+		ret = edgetpu_ioctl_sync_fence_create(client, argp);
 		break;
 	case EDGETPU_SIGNAL_SYNC_FENCE:
 		ret = edgetpu_ioctl_sync_fence_signal(argp);
@@ -1023,6 +1030,20 @@ static const struct file_operations mappings_ops = {
 	.release = single_release,
 };
 
+static int edgetpu_pm_debugfs_set_wakelock(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = data;
+	int ret = 0;
+
+	if (val)
+		ret = edgetpu_pm_get(etdev->pm);
+	else
+		edgetpu_pm_put(etdev->pm);
+	return ret;
+}
+DEFINE_DEBUGFS_ATTRIBUTE(fops_wakelock, NULL, edgetpu_pm_debugfs_set_wakelock,
+			 "%llu\n");
+
 static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 {
 	etdev->d_entry =
@@ -1033,6 +1054,8 @@ static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 	}
 	debugfs_create_file("mappings", 0440, etdev->d_entry,
 			    etdev, &mappings_ops);
+	debugfs_create_file("wakelock", 0220, etdev->d_entry, etdev,
+			    &fops_wakelock);
 #ifndef EDGETPU_FEATURE_MOBILE
 	debugfs_create_file("statusregs", 0440, etdev->d_entry, etdev,
 			    &statusregs_ops);
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index 95578b5..12b7909 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -6,6 +6,7 @@
  */
 
 #include <linux/device.h>
+#include <linux/dma-iommu.h>
 #include <linux/dma-mapping.h>
 #include <linux/idr.h>
 #include <linux/iommu.h>
@@ -243,6 +244,14 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 		etdev_warn(etdev, "AUX domains not supported\n");
 	else
 		etiommu->aux_enabled = true;
+
+#if IS_ENABLED(CONFIG_ANDROID)
+	/* Enable best fit algorithm to minimize fragmentation */
+	ret = iommu_dma_enable_best_fit_algo(etdev->dev);
+	if (ret)
+		etdev_warn(etdev, "Failed to enable best-fit IOVA allocator (%d)\n", ret);
+#endif
+
 	ret = check_default_domain(etdev, etiommu);
 	if (ret)
 		goto err_free;
@@ -349,18 +358,12 @@ int edgetpu_mmu_map(struct edgetpu_dev *etdev, struct edgetpu_mapping *map,
 		iommu_get_domain_for_dev(etdev->dev);
 
 	ret = get_iommu_map_params(etdev, map, context_id, &params, mmu_flags);
-
 	if (ret)
 		return ret;
 
-	if (mmu_flags & EDGETPU_MMU_64)
-		dev_warn_once(etdev->dev,
-			      "%s: 64-bit addressing is not supported",
-			      __func__);
-
 	ret = dma_map_sg_attrs(etdev->dev, map->sgt.sgl, map->sgt.nents, map->dir, map->dma_attrs);
 	if (!ret)
-		return -EINVAL;
+		return -ENOSPC;
 	map->sgt.nents = ret;
 	iova = sg_dma_address(map->sgt.sgl);
 
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 834318f..0c6f5ad 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -38,8 +38,8 @@
 /* fake-firmware could respond in a short time */
 #define KCI_TIMEOUT	(200)
 #else
-/* 5 secs. */
-#define KCI_TIMEOUT	(5000)
+/* Wait for up to 1 second for FW to respond. */
+#define KCI_TIMEOUT	(1000)
 #endif
 
 /* A macro for KCIs to leave early when the device state is known to be bad. */
@@ -246,10 +246,9 @@ edgetpu_kci_handle_response(struct edgetpu_kci *kci,
 		int ret = edgetpu_reverse_kci_add_response(kci, resp);
 
 		if (ret)
-			etdev_warn(
-				kci->mailbox->etdev,
-				"Failed to handle reverse KCI code %u (%d)\n",
-				resp->code, ret);
+			etdev_warn_ratelimited(kci->mailbox->etdev,
+					       "Failed to handle reverse KCI code %u (%d)\n",
+					       resp->code, ret);
 		return;
 	}
 	/*
@@ -642,7 +641,12 @@ int edgetpu_kci_push_cmd(struct edgetpu_kci *kci,
 
 	mutex_lock(&kci->cmd_queue_lock);
 
-	cmd->seq = kci->cur_seq;
+	/*
+	 * Only update sequence number for KCI commands.  Do not change
+	 * sequence number for responses to RKCI commands.
+	 */
+	if (!(cmd->seq & KCI_REVERSE_FLAG))
+		cmd->seq = kci->cur_seq;
 	/*
 	 * The lock ensures mailbox->cmd_queue_tail cannot be changed by
 	 * other processes (this method should be the only one to modify the
@@ -681,7 +685,8 @@ int edgetpu_kci_push_cmd(struct edgetpu_kci *kci,
 	/* triggers doorbell */
 	EDGETPU_MAILBOX_CMD_QUEUE_WRITE_SYNC(kci->mailbox, doorbell_set, 1);
 	/* bumps sequence number after the command is sent */
-	kci->cur_seq++;
+	if (!(cmd->seq & KCI_REVERSE_FLAG))
+		kci->cur_seq++;
 	ret = 0;
 out:
 	mutex_unlock(&kci->cmd_queue_lock);
@@ -707,6 +712,10 @@ static int edgetpu_kci_send_cmd_return_resp(
 	ret = edgetpu_kci_push_cmd(kci, cmd, resp);
 	if (ret)
 		return ret;
+
+	if (!resp)
+		return 0;
+
 	ret = wait_event_timeout(kci->wait_list_waitq,
 				 resp->status != KCI_STATUS_WAITING_RESPONSE,
 				 msecs_to_jiffies(KCI_TIMEOUT));
@@ -755,22 +764,11 @@ int edgetpu_kci_send_cmd(struct edgetpu_kci *kci,
 {
 	struct edgetpu_kci_response_element resp;
 
-	return edgetpu_kci_send_cmd_return_resp(kci, cmd, &resp);
-}
-
-int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
-			     u32 size, enum dma_data_direction dir)
-{
-	struct edgetpu_command_element cmd = {
-		.code = KCI_CODE_UNMAP_BUFFER,
-		.dma = {
-			.address = tpu_addr,
-			.size = size,
-			.flags = dir,
-		},
-	};
-
-	return edgetpu_kci_send_cmd(kci, &cmd);
+	/* Don't wait on a response for reverse KCI response. */
+	if (cmd->seq & KCI_REVERSE_FLAG)
+		return edgetpu_kci_send_cmd_return_resp(kci, cmd, NULL);
+	else
+		return edgetpu_kci_send_cmd_return_resp(kci, cmd, &resp);
 }
 
 int edgetpu_kci_map_log_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
@@ -926,6 +924,9 @@ int edgetpu_kci_update_usage(struct edgetpu_dev *etdev)
 
 fw_unlock:
 	edgetpu_firmware_unlock(etdev);
+
+	if (ret)
+		etdev_warn_once(etdev, "get firmware usage stats failed: %d", ret);
 	return ret;
 }
 
@@ -1091,3 +1092,17 @@ int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block)
 
 	return edgetpu_kci_send_cmd(etdev->kci, &cmd);
 }
+
+int edgetpu_kci_resp_rkci_ack(struct edgetpu_dev *etdev,
+			      struct edgetpu_kci_response_element *rkci_cmd)
+{
+	struct edgetpu_command_element cmd = {
+		.seq = rkci_cmd->seq,
+		.code = KCI_CODE_RKCI_ACK,
+	};
+
+	if (!etdev->kci)
+		return -ENODEV;
+
+	return edgetpu_kci_send_cmd(etdev->kci, &cmd);
+}
diff --git a/drivers/edgetpu/edgetpu-kci.h b/drivers/edgetpu/edgetpu-kci.h
index cc8bff3..a11a181 100644
--- a/drivers/edgetpu/edgetpu-kci.h
+++ b/drivers/edgetpu/edgetpu-kci.h
@@ -60,7 +60,11 @@ struct edgetpu_dma_descriptor {
 };
 
 struct edgetpu_command_element {
-	u64 seq;	 /* set by edgetpu_kci_push_cmd() */
+	/*
+	 * Set by edgetpu_kci_push_cmd() in case of KCI cmd and copied from
+	 * the RKCI cmd in case of RKCI response.
+	 */
+	u64 seq;
 	u16 code;
 	u16 reserved[3]; /* explicit padding, does not affect alignment */
 	struct edgetpu_dma_descriptor dma;
@@ -72,7 +76,8 @@ struct edgetpu_kci_response_element {
 	/*
 	 * Reserved for host use - firmware can't touch this.
 	 * If a value is written here it will be discarded and overwritten
-	 * during response processing.
+	 * during response processing. However, when repurposed as an RKCI
+	 * command, the FW can set this field.
 	 */
 	u16 status;
 	/*
@@ -110,6 +115,8 @@ enum edgetpu_kci_code {
 	KCI_CODE_GET_USAGE = 12,
 	KCI_CODE_NOTIFY_THROTTLING = 13,
 	KCI_CODE_BLOCK_BUS_SPEED_CONTROL = 14,
+
+	KCI_CODE_RKCI_ACK = 256,
 };
 
 /*
@@ -284,14 +291,6 @@ int edgetpu_kci_send_cmd(struct edgetpu_kci *kci,
 			 struct edgetpu_command_element *cmd);
 
 /*
- * Sends the "Unmap Buffer Sync" command and waits for remote response.
- *
- * Returns the code of response, or a negative errno on error.
- */
-int edgetpu_kci_unmap_buffer(struct edgetpu_kci *kci, tpu_addr_t tpu_addr,
-			     u32 size, enum dma_data_direction dir);
-
-/*
  * Sends a FIRMWARE_INFO command and expects a response with a
  * edgetpu_fw_info struct filled out, including what firmware type is running,
  * along with build CL and time.
@@ -405,4 +404,14 @@ int edgetpu_kci_notify_throttling(struct edgetpu_dev *etdev, u32 level);
  */
 int edgetpu_kci_block_bus_speed_control(struct edgetpu_dev *etdev, bool block);
 
+/*
+ * Send an ack to the FW after handling a reverse KCI request.
+ *
+ * The FW may wait for a response from the kernel for an RKCI request so a
+ * response could be sent as an ack.
+ */
+int edgetpu_kci_resp_rkci_ack(struct edgetpu_dev *etdev,
+			      struct edgetpu_kci_response_element *rkci_cmd);
+
+
 #endif /* __EDGETPU_KCI_H__ */
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index af6bcb7..96fe245 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -22,6 +22,8 @@
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
 
+static struct edgetpu_dev *edgetpu_debug_pointer;
+
 /*
  * Log and trace buffers at the beginning of the remapped region,
  * pool memory afterwards.
@@ -220,6 +222,20 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
+static void edgetpu_platform_parse_pmu(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct device *dev = etdev->dev;
+	u32 reg;
+
+	if (of_find_property(dev->of_node, "pmu-status-base", NULL) &&
+	    !of_property_read_u32_index(dev->of_node, "pmu-status-base", 0, &reg)) {
+		etmdev->pmu_status = devm_ioremap(dev, reg, 0x4);
+		if (!etmdev->pmu_status)
+			etdev_info(etdev, "Using ACPM for blk status query\n");
+	}
+}
+
 static int edgetpu_platform_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
 {
 	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
@@ -397,6 +413,8 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 	if (ret)
 		dev_warn(dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
 
+	edgetpu_platform_parse_pmu(etmdev);
+
 	etmdev->log_mem = devm_kcalloc(dev, etdev->num_cores, sizeof(*etmdev->log_mem), GFP_KERNEL);
 	if (!etmdev->log_mem) {
 		ret = -ENOMEM;
@@ -438,6 +456,8 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 	/* Turn the device off unless a client request is already received. */
 	edgetpu_pm_shutdown(etdev, false);
 
+	edgetpu_debug_pointer = etdev;
+
 	return 0;
 out_destroy_fw:
 	edgetpu_mobile_firmware_destroy(etdev);
@@ -473,6 +493,9 @@ static int edgetpu_mobile_platform_remove(struct platform_device *pdev)
 	edgetpu_platform_cleanup_fw_region(etmdev);
 	edgetpu_pm_put(etdev->pm);
 	edgetpu_pm_shutdown(etdev, true);
-	mobile_pm_destroy(etdev);
+	edgetpu_mobile_pm_destroy(etdev);
+
+	edgetpu_debug_pointer = NULL;
+
 	return 0;
 }
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index 9d41571..335eb94 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -45,8 +45,8 @@ struct edgetpu_mobile_platform_pwr {
 	int (*lpm_up)(struct edgetpu_dev *etdev);
 	void (*lpm_down)(struct edgetpu_dev *etdev);
 
-	/* Block shutdown callback, may be NULL */
-	void (*block_down)(struct edgetpu_dev *etdev);
+	/* Block shutdown status callback, may be NULL */
+	bool (*is_block_down)(struct edgetpu_dev *etdev);
 
 	/* Firmware shutdown callback. Must be implemented */
 	void (*firmware_down)(struct edgetpu_dev *etdev);
@@ -105,6 +105,8 @@ struct edgetpu_mobile_platform_dev {
 	int n_irq;
 	/* Array of IRQ numbers */
 	int *irq;
+	/* PMU status base address for block status, maybe NULL */
+	void __iomem *pmu_status;
 
 	/* callbacks for chip-dependent implementations */
 
diff --git a/drivers/edgetpu/mobile-firmware.c b/drivers/edgetpu/mobile-firmware.c
index c1d58d6..42f982f 100644
--- a/drivers/edgetpu/mobile-firmware.c
+++ b/drivers/edgetpu/mobile-firmware.c
@@ -281,35 +281,6 @@ static int mobile_firmware_gsa_authenticate(struct edgetpu_mobile_platform_dev *
 	return ret;
 }
 
-/* TODO(b/197074886): remove once rio has GSA support */
-static void program_iremap_csr(struct edgetpu_dev *etdev)
-{
-	const int ctx_id = 0, sid0 = 0x30, sid1 = 0x34;
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	phys_addr_t fw_paddr = etmdev->fw_region_paddr;
-
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY, (ctx_id << 16) | sid0);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY + 8,
-			     (ctx_id << 16) | sid1);
-#if defined(ZEBU_SYSMMU_WORKAROUND)
-	/*
-	 * This is required on ZeBu after b/197718405 is fixed, which forwards all transactions to
-	 * the non-secure SysMMU.
-	 */
-	fw_paddr = EDGETPU_INSTRUCTION_REMAP_BASE;
-#endif
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE, fw_paddr);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE + 8, fw_paddr);
-
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT,
-			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_16M);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT + 8,
-			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_16M);
-
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL, 1);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL + 8, 1);
-}
-
 static void mobile_firmware_setup_ssmt(struct edgetpu_dev *etdev)
 {
 	int i;
@@ -331,8 +302,6 @@ static void mobile_firmware_setup_ssmt(struct edgetpu_dev *etdev)
 	 * Reset the table to zeroes if running non-secure firmware, since the SSMT
 	 * will be in clamped mode and we want all memory accesses to go to the
 	 * default page table.
-	 *
-	 * TODO(b/204384254) Confirm SSMT setup for Rio
 	 */
 	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
 		int val;
@@ -355,9 +324,6 @@ static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
 	edgetpu_mailbox_reset(etdev->kci->mailbox);
 
-	if (IS_ENABLED(CONFIG_RIO))
-		program_iremap_csr(etdev);
-
 	mobile_firmware_setup_ssmt(etdev);
 
 	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
@@ -390,6 +356,7 @@ static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	phys_addr_t image_start, image_end, carveout_start, carveout_end;
 	bool image_config_changed;
+	struct mobile_image_header *hdr;
 
 	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
 		etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
@@ -404,6 +371,11 @@ static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 		return -ENOMEM;
 	}
 
+	hdr = (struct mobile_image_header *)fw_buf->vaddr;
+	if (hdr->Magic != EDGETPU_MOBILE_FW_MAGIC)
+		etdev_warn(etdev, "Invalid firmware header magic value %#08x\n",
+			   hdr->Magic);
+
 	/* fetch the firmware versions */
 	image_config = fw_buf->vaddr + MOBILE_IMAGE_CONFIG_OFFSET;
 	memcpy(&etdev->fw_version, &image_config->firmware_versions,
diff --git a/drivers/edgetpu/mobile-firmware.h b/drivers/edgetpu/mobile-firmware.h
index 27ee048..05102c9 100644
--- a/drivers/edgetpu/mobile-firmware.h
+++ b/drivers/edgetpu/mobile-firmware.h
@@ -75,6 +75,9 @@ struct mobile_image_header {
 	struct mobile_image_config ImageConfig;
 };
 
+/* Value of Magic field above: 'TPUF' as a 32-bit LE int */
+#define EDGETPU_MOBILE_FW_MAGIC	0x46555054
+
 int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev);
 void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev);
 
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 52a43ab..50c3866 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -64,6 +64,7 @@ static int mobile_pwr_state_init(struct device *dev)
 		ret = pm_runtime_get_sync(dev);
 		if (ret) {
 			pm_runtime_put_noidle(dev);
+			pm_runtime_disable(dev);
 			dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
 			return ret;
 		}
@@ -74,6 +75,7 @@ static int mobile_pwr_state_init(struct device *dev)
 		dev_err(dev, "error initializing tpu state: %d\n", ret);
 		if (curr_state > TPU_OFF)
 			pm_runtime_put_sync(dev);
+		pm_runtime_disable(dev);
 		return ret;
 	}
 
@@ -255,8 +257,6 @@ static int mobile_pwr_state_set_locked(struct edgetpu_mobile_platform_dev *etmde
 			dev_err(dev, "%s: pm_runtime_put_sync returned %d\n", __func__, ret);
 			return ret;
 		}
-		if (platform_pwr->block_down)
-			platform_pwr->block_down(etdev);
 	}
 
 	return ret;
@@ -432,7 +432,12 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 	struct edgetpu_dev *etdev = etpm->etdev;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
-	int ret = mobile_pwr_state_set(etpm->etdev, mobile_get_initial_pwr_state(etdev->dev));
+	int ret;
+
+	if (platform_pwr->is_block_down && !platform_pwr->is_block_down(etdev))
+		return -EAGAIN;
+
+	ret = mobile_pwr_state_set(etpm->etdev, mobile_get_initial_pwr_state(etdev->dev));
 
 	etdev_info(etpm->etdev, "Powering up\n");
 
@@ -600,6 +605,7 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct device *dev = etdev->dev;
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int curr_state;
 
 	ret = mobile_pwr_state_init(dev);
 	if (ret)
@@ -641,7 +647,19 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 
 	if (platform_pwr->after_create)
 		ret = platform_pwr->after_create(etdev);
+	if (ret)
+		goto err_debugfs_remove;
 
+	return 0;
+
+err_debugfs_remove:
+	debugfs_remove_recursive(platform_pwr->debugfs_dir);
+	/* pm_runtime_{enable,get_sync} were called in mobile_pwr_state_init */
+
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	if (curr_state > TPU_OFF)
+		pm_runtime_put_sync(dev);
+	pm_runtime_disable(dev);
 	return ret;
 }
 
@@ -668,17 +686,17 @@ static struct edgetpu_pm_handlers mobile_pm_handlers = {
 	.power_down = mobile_power_down,
 };
 
-int mobile_pm_create(struct edgetpu_dev *etdev)
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev)
 {
 	return edgetpu_pm_create(etdev, &mobile_pm_handlers);
 }
 
-void mobile_pm_destroy(struct edgetpu_dev *etdev)
+void edgetpu_mobile_pm_destroy(struct edgetpu_dev *etdev)
 {
 	edgetpu_pm_destroy(etdev);
 }
 
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
+void edgetpu_mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
 {
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
@@ -741,7 +759,7 @@ static void mobile_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 	mutex_unlock(&platform_pwr->scenario_lock);
 }
 
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
+void edgetpu_mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
 {
 	etdev_dbg(etdev, "%s: bts request - val = %u\n", __func__, bts_val);
 
diff --git a/drivers/edgetpu/mobile-pm.h b/drivers/edgetpu/mobile-pm.h
index 4299b49..b0f8c3d 100644
--- a/drivers/edgetpu/mobile-pm.h
+++ b/drivers/edgetpu/mobile-pm.h
@@ -47,6 +47,8 @@ static inline int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
 enum mobile_reverse_kci_code {
 	RKCI_CODE_PM_QOS = RKCI_CHIP_CODE_FIRST + 1,
 	RKCI_CODE_BTS = RKCI_CHIP_CODE_FIRST + 2,
+	/* The above codes have been deprecated. */
+
 	RKCI_CODE_PM_QOS_BTS = RKCI_CHIP_CODE_FIRST + 3,
 };
 
@@ -68,7 +70,7 @@ enum mobile_reverse_kci_code {
  * chipsets.
  * Needs to be called after the devices's platform_pwr struct has been initialized.
  */
-int mobile_pm_create(struct edgetpu_dev *etdev);
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev);
 
 /*
  * Wrapper for chip-specific implementation.
@@ -79,12 +81,12 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev);
 /*
  * Destroy power management interface for an edgetpu device on mobile chipsets.
  */
-void mobile_pm_destroy(struct edgetpu_dev *etdev);
+void edgetpu_mobile_pm_destroy(struct edgetpu_dev *etdev);
 
 /* Set required QoS value for the edgetpu device. */
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
+void edgetpu_mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
 
 /* Set BTS value for the edgetpu device. */
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val);
+void edgetpu_mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val);
 
 #endif /* __MOBILE_PM_H__ */