Revert "[Copybara Auto Merge] Merge branch 'pro' into android-gs-cloudripper-5.10"

This reverts commit 75813437d78bd26b5a547c0b6bedbfe26d4c0664.

Reason for revert: b/197285064 - inference output mismatches. It's not yet merged into googleplex-android. Reverting until the root cause of the bug is figured out.

Change-Id: I99f19f08fb53a1dafe270a07c271c03fe2a46b13

28 files changed, 600 insertions, 328 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 2caadf4..170a69e 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -13,8 +13,6 @@ 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-mobile-objs	:= edgetpu-mobile-firmware.o
-edgetpu-mcp-objs	:= edgetpu-shared-fw.o
-
 
 janeiro-y	:= janeiro-device.o janeiro-device-group.o janeiro-fs.o janeiro-core.o janeiro-platform.o janeiro-firmware.o janeiro-pm.o janeiro-debug-dump.o janeiro-usage-stats.o janeiro-iommu.o janeiro-wakelock.o $(edgetpu-mobile-objs) $(edgetpu-objs)
 
diff --git a/drivers/edgetpu/Kconfig b/drivers/edgetpu/Kconfig
index 4970c54..0436086 100644
--- a/drivers/edgetpu/Kconfig
+++ b/drivers/edgetpu/Kconfig
@@ -10,6 +10,17 @@ config EDGETPU_FRAMEWORK
 	help
 	  This framework supports darwinn-2.0 devices.
 
+config ABROLHOS
+	tristate "Abrolhos ML accelerator direct host control driver"
+	depends on EDGETPU_FRAMEWORK
+	select PM
+	help
+	  This driver supports the Abrolhos device.  Say Y if you want to
+	  include this driver in the kernel.
+
+	  To compile this driver as a module, choose M here.  The module
+	  will be called "abrolhos".
+
 config JANEIRO
 	tristate "Janeiro ML accelerator device driver"
 	depends on EDGETPU_FRAMEWORK
@@ -21,6 +32,24 @@ config JANEIRO
 	  To compile this driver as a module, choose M here.  The module
 	  will be called "janeiro".
 
+config HERMOSA
+	tristate "Hermosa edgetpu ML accelerator direct host control driver"
+	depends on EDGETPU_FRAMEWORK
+	select INTERVAL_TREE if EDGETPU_DEVICE_DRAM
+	help
+	  This driver supports the Hermosa device.  Say Y if you want to
+	  include this driver in the kernel.
+
+	  To compile this driver as a module, choose M here.  The module
+	  will be called "hermosa".
+
+config EDGETPU_DEVICE_DRAM
+	bool "EdgeTPU Device DRAM support"
+	depends on HERMOSA
+	default y
+	help
+	  Say Y if the target device has the embedded device DRAM.
+
 config EDGETPU_EXTERNAL_WRAPPER_CLASS
         bool "EdgeTPU add external wrapper class"
 	depends on EDGETPU_FRAMEWORK
@@ -43,4 +72,6 @@ config EDGETPU_TELEMETRY_TRACE
 	  It's fine to have this enabled even the firmware doesn't send tracing
 	  events.
 
+source "drivers/edgetpu/unittests/Kconfig"
+
 endmenu
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index f027fa6..caf01c0 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -19,8 +19,6 @@ edgetpu-objs	:= edgetpu-async.o edgetpu-dmabuf.o edgetpu-iremap-pool.o \
 		   edgetpu-firmware-util.o edgetpu-firmware.o
 
 edgetpu-mobile-objs	:= edgetpu-mobile-firmware.o
-edgetpu-mcp-objs	:= edgetpu-shared-fw.o
-
 
 janeiro-objs	:= janeiro-core.o janeiro-debug-dump.o janeiro-device-group.o \
 		   janeiro-device.o janeiro-firmware.o janeiro-fs.o \
@@ -28,7 +26,5 @@ janeiro-objs	:= janeiro-core.o janeiro-debug-dump.o janeiro-device-group.o \
 		   janeiro-usage-stats.o janeiro-wakelock.o \
 		   $(edgetpu-objs) $(edgetpu-mobile-objs)
 
-KBUILD_OPTIONS += CONFIG_JANEIRO=m
-
 modules modules_install clean:
 	$(MAKE) -C $(KERNEL_SRC) M=$(M) W=1 $(KBUILD_OPTIONS) $(@)
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index 9306fac..5a13adb 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -8,10 +8,28 @@
 #ifndef __EDGETPU_CONFIG_H__
 #define __EDGETPU_CONFIG_H__
 
+#ifdef CONFIG_HERMOSA
+
+#include "hermosa/config.h"
+
+#else /* !CONFIG_HERMOSA */
+
 #ifdef CONFIG_JANEIRO
 
 #include "janeiro/config.h"
+
+#else
+
+#ifndef CONFIG_ABROLHOS
+#define CONFIG_ABROLHOS
+#warning "Building default chipset abrolhos"
+#endif
+
+#include "abrolhos/config.h"
+
 #endif /* CONFIG_JANEIRO */
+#endif /* CONFIG_HERMOSA */
+
 #define EDGETPU_DEFAULT_FIRMWARE_NAME "google/edgetpu-" DRIVER_NAME ".fw"
 #define EDGETPU_TEST_FIRMWARE_NAME "google/edgetpu-" DRIVER_NAME "-test.fw"
 
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index e21ac3d..3115f17 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -545,8 +545,6 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	 */
 	if (client->group)
 		edgetpu_device_group_leave(client);
-	/* invoke chip-dependent removal handler before releasing resources */
-	edgetpu_chip_client_remove(client);
 	edgetpu_wakelock_free(client->wakelock);
 	/*
 	 * It should be impossible to access client->wakelock after this cleanup
@@ -563,6 +561,7 @@ void edgetpu_client_remove(struct edgetpu_client *client)
 	    1 << perdie_event_id_to_num(EDGETPU_PERDIE_EVENT_TRACES_AVAILABLE))
 		edgetpu_telemetry_unset_event(etdev, EDGETPU_TELEMETRY_TRACE);
 
+	edgetpu_chip_client_remove(client);
 	edgetpu_client_put(client);
 }
 
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 402d231..d1c0832 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -41,6 +41,9 @@
 #include "edgetpu-p2p-mailbox.h"
 #endif
 
+#define for_each_list_group_client(c, group) \
+	list_for_each_entry(c, &group->clients, list)
+
 #define for_each_list_group_client_safe(c, n, group) \
 	list_for_each_entry_safe(c, n, &group->clients, list)
 
@@ -140,11 +143,16 @@ static int edgetpu_group_activate(struct edgetpu_device_group *group)
 static void edgetpu_group_deactivate(struct edgetpu_device_group *group)
 {
 	u8 mailbox_id;
+	int ret;
 
 	if (edgetpu_group_mailbox_detached_locked(group))
 		return;
 	mailbox_id = edgetpu_group_context_id_locked(group);
-	edgetpu_mailbox_deactivate(group->etdev, mailbox_id);
+	ret = edgetpu_mailbox_deactivate(group->etdev, mailbox_id);
+	if (ret)
+		etdev_err(group->etdev, "deactivate mailbox for VCID %d failed with %d",
+			  group->vcid, ret);
+	return;
 }
 
 /*
@@ -1115,12 +1123,6 @@ static void edgetpu_host_map_show(struct edgetpu_mapping *map,
 	}
 }
 
-size_t edgetpu_group_mappings_total_size(struct edgetpu_device_group *group)
-{
-	return edgetpu_mappings_total_size(&group->host_mappings) +
-		edgetpu_mappings_total_size(&group->dmabuf_mappings);
-}
-
 /*
  * Pins the user-space address @arg->host_address and returns the pinned pages.
  * @pnum_pages is set to the number of pages.
@@ -1140,7 +1142,6 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	int i;
 	int ret;
 	struct vm_area_struct *vma;
-	struct vm_area_struct **vmas;
 	unsigned int foll_flags = FOLL_LONGTERM | FOLL_WRITE;
 
 	if (size == 0)
@@ -1157,10 +1158,11 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	 */
 	pages = kvmalloc((num_pages * sizeof(*pages)), GFP_KERNEL | __GFP_NOWARN);
 	if (!pages) {
-		etdev_dbg(etdev, "%s: kvmalloc pages failed (%lu bytes)\n",
-			  __func__, (num_pages * sizeof(*pages)));
+		etdev_dbg(etdev, "%s: kvmalloc failed (%lu bytes)\n", __func__,
+			  (num_pages * sizeof(*pages)));
 		return ERR_PTR(-ENOMEM);
 	}
+
 	/*
 	 * The host pages might be read-only and could fail if we attempt to pin
 	 * it with FOLL_WRITE.
@@ -1174,43 +1176,10 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 		*preadonly = false;
 	}
 
-	/* Try fast call first, in case it's actually faster. */
 	ret = pin_user_pages_fast(host_addr & PAGE_MASK, num_pages, foll_flags,
 				  pages);
-	if (ret == num_pages) {
-		*pnum_pages = num_pages;
-		return pages;
-	}
-	if (ret < 0) {
-		etdev_dbg(etdev, "pin_user_pages failed %u:%pK-%u: %d",
-			  group->workload_id, (void *)host_addr, num_pages,
-			  ret);
-		if (ret != -ENOMEM) {
-			num_pages = 0;
-			goto error;
-		}
-	}
-	etdev_dbg(etdev,
-		  "pin_user_pages_fast error %u:%pK npages=%u ret=%d",
-		  group->workload_id, (void *)host_addr, num_pages,
-		  ret);
-	/* Unpin any partial mapping and start over again. */
-	for (i = 0; i < ret; i++)
-		unpin_user_page(pages[i]);
-
-	/* Allocate our own vmas array non-contiguous. */
-	vmas = kvmalloc((num_pages * sizeof(*vmas)), GFP_KERNEL | __GFP_NOWARN);
-	if (!vmas) {
-		etdev_dbg(etdev, "%s: kvmalloc vmas failed (%lu bytes)\n",
-			  __func__, (num_pages * sizeof(*pages)));
-		kvfree(pages);
-		return ERR_PTR(-ENOMEM);
-	}
-	ret = pin_user_pages(host_addr & PAGE_MASK, num_pages, foll_flags,
-			     pages, vmas);
-	kvfree(vmas);
 	if (ret < 0) {
-		etdev_dbg(etdev, "pin_user_pages failed %u:%pK-%u: %d",
+		etdev_dbg(etdev, "get user pages failed %u:%pK-%u: %d",
 			  group->workload_id, (void *)host_addr, num_pages,
 			  ret);
 		num_pages = 0;
@@ -1218,7 +1187,7 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	}
 	if (ret < num_pages) {
 		etdev_dbg(etdev,
-			  "pin_user_pages partial %u:%pK npages=%u pinned=%d",
+			  "get user pages partial %u:%pK npages=%u pinned=%d",
 			  group->workload_id, (void *)host_addr, num_pages,
 			  ret);
 		num_pages = ret;
@@ -1227,6 +1196,7 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	}
 
 	*pnum_pages = num_pages;
+
 	return pages;
 
 error:
@@ -1487,7 +1457,6 @@ int edgetpu_device_group_map(struct edgetpu_device_group *group,
 			goto error;
 	}
 
-	map->map_size = arg->size;
 	ret = edgetpu_mapping_add(&group->host_mappings, map);
 	if (ret) {
 		etdev_dbg(etdev, "duplicate mapping %u:0x%llx",
@@ -1608,10 +1577,7 @@ void edgetpu_mappings_clear_group(struct edgetpu_device_group *group)
 void edgetpu_group_mappings_show(struct edgetpu_device_group *group,
 				 struct seq_file *s)
 {
-	enum edgetpu_context_id context =
-		edgetpu_group_context_id_locked(group);
-
-	seq_printf(s, "group %u", group->workload_id);
+	seq_printf(s, "workload %u", group->workload_id);
 	switch (group->status) {
 	case EDGETPU_DEVICE_GROUP_WAITING:
 	case EDGETPU_DEVICE_GROUP_FINALIZED:
@@ -1623,23 +1589,14 @@ void edgetpu_group_mappings_show(struct edgetpu_device_group *group,
 		seq_puts(s, ": disbanded\n");
 		return;
 	}
-
-	if (context == EDGETPU_CONTEXT_INVALID)
-		seq_puts(s, " context (none):\n");
-	else if (context & EDGETPU_CONTEXT_DOMAIN_TOKEN)
-		seq_printf(s, " context detached 0x%x:\n",
-			   context & ~(EDGETPU_CONTEXT_DOMAIN_TOKEN));
-	else
-		seq_printf(s, " context mbox %d:\n", context);
+	seq_printf(s, " context %d:\n", edgetpu_group_context_id_locked(group));
 
 	if (group->host_mappings.count) {
-		seq_printf(s, "host buffer mappings (%zd):\n",
-			   group->host_mappings.count);
+		seq_puts(s, "host buffer mappings:\n");
 		edgetpu_mappings_show(&group->host_mappings, s);
 	}
 	if (group->dmabuf_mappings.count) {
-		seq_printf(s, "dma-buf buffer mappings (%zd):\n",
-			   group->dmabuf_mappings.count);
+		seq_puts(s, "dma-buf buffer mappings:\n");
 		edgetpu_mappings_show(&group->dmabuf_mappings, s);
 	}
 
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 5c5eaaf..7ec262f 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -157,11 +157,6 @@ struct edgetpu_list_group {
 	for (l = list_entry(etdev->groups.next, typeof(*l), list), g = l->grp; \
 	     &l->list != &etdev->groups;                                       \
 	     l = list_entry(l->list.next, typeof(*l), list), g = l->grp)
-
-/* Loop through group->clients (hold group->lock prior). */
-#define for_each_list_group_client(c, group) \
-	list_for_each_entry(c, &group->clients, list)
-
 /*
  * Returns if the group is waiting for members to join.
  *
@@ -335,9 +330,6 @@ int edgetpu_device_group_sync_buffer(struct edgetpu_device_group *group,
 /* Clear all mappings for a device group. */
 void edgetpu_mappings_clear_group(struct edgetpu_device_group *group);
 
-/* Return total size of all mappings for the group in bytes */
-size_t edgetpu_group_mappings_total_size(struct edgetpu_device_group *group);
-
 /*
  * Return context ID for group MMU mappings.
  *
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index 2148f1f..b0dc9a1 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -670,7 +670,7 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 
 	get_dma_buf(dmabuf);
 	dmap->dmabufs[0] = dmabuf;
-	dmap->map.map_size = dmap->size = size = dmabuf->size;
+	dmap->size = size = dmabuf->size;
 	if (IS_MIRRORED(flags)) {
 		for (i = 0; i < group->n_clients; i++) {
 			etdev = edgetpu_device_group_nth_etdev(group, i);
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index f24135b..1cdfaa1 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -285,20 +285,4 @@ uint32_t edgetpu_firmware_get_cl(struct edgetpu_firmware *et_fw);
 /* Returns the build time of the image in seconds since 1970. */
 uint64_t edgetpu_firmware_get_build_time(struct edgetpu_firmware *et_fw);
 
-/*
- * Kernel verify firmware signature (if EDGETPU_FEATURE_FW_SIG enabled).
- *
- * @etdev:      the edgetpu_dev for which the initial load of a (probably
- *              shared) firmware image is requested
- * @name:       name of the image being validated (request_firmware path)
- * @image_data: passes in the pointer to the raw image with signature, returns
- *              pointer to the firmware code image.
- * @image_size: passes in the size of the raw image with signature, returns
- *              size of the firmware code image.
- */
-bool edgetpu_firmware_verify_signature(struct edgetpu_dev *etdev,
-				       const char *name,
-				       void **image_data, size_t *image_size);
-
-
 #endif /* __EDGETPU_FIRMWARE_H__ */
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 3598104..0dcad64 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -1013,122 +1013,26 @@ static ssize_t clients_show(
 {
 	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
 	struct edgetpu_list_device_client *lc;
-	ssize_t len;
 	ssize_t ret = 0;
 
 	mutex_lock(&etdev->clients_lock);
 	for_each_list_device_client(etdev, lc) {
-		struct edgetpu_device_group *group;
-
-		mutex_lock(&lc->client->group_lock);
-		group = lc->client->group;
-		len = scnprintf(buf, PAGE_SIZE - ret,
-				"pid %d tgid %d group %d wakelock %d\n",
-				lc->client->pid, lc->client->tgid,
-				group ? group->workload_id : -1,
-				NO_WAKELOCK(lc->client->wakelock) ?
-				0 : lc->client->wakelock->req_count);
-		mutex_unlock(&lc->client->group_lock);
-		buf += len;
-		ret += len;
+		ret += scnprintf(buf, PAGE_SIZE - ret,
+				 "pid %d tgid %d wakelock %d\n",
+				 lc->client->pid, lc->client->tgid,
+				 NO_WAKELOCK(lc->client->wakelock) ?
+				 0 : lc->client->wakelock->req_count);
+		buf += ret;
 	}
 	mutex_unlock(&etdev->clients_lock);
 	return ret;
 }
 static DEVICE_ATTR_RO(clients);
 
-static ssize_t show_group(struct edgetpu_dev *etdev,
-			  struct edgetpu_device_group *group, char *buf,
-			  ssize_t buflen)
-{
-	enum edgetpu_context_id context =
-		edgetpu_group_context_id_locked(group);
-	struct edgetpu_list_group_client *lc;
-	ssize_t len;
-	ssize_t ret = 0;
-
-	len = scnprintf(buf, buflen - ret, "group %u ", group->workload_id);
-	buf += len;
-	ret += len;
-
-	switch (group->status) {
-	case EDGETPU_DEVICE_GROUP_WAITING:
-		len = scnprintf(buf, buflen - ret, "forming ");
-		buf += len;
-		ret += len;
-		break;
-	case EDGETPU_DEVICE_GROUP_FINALIZED:
-		break;
-	case EDGETPU_DEVICE_GROUP_ERRORED:
-		len = scnprintf(buf, buflen - ret, "error 0x%x ",
-				group->fatal_errors);
-		buf += len;
-		ret += len;
-		break;
-	case EDGETPU_DEVICE_GROUP_DISBANDED:
-		len = scnprintf(buf, buflen - ret, "disbanded\n");
-		ret += len;
-		return ret;
-	}
-
-	if (context == EDGETPU_CONTEXT_INVALID)
-		len = scnprintf(buf, buflen - ret, "context (none) ");
-	else if (context & EDGETPU_CONTEXT_DOMAIN_TOKEN)
-		len = scnprintf(buf, buflen - ret, "context detached 0x%x ",
-				context & ~(EDGETPU_CONTEXT_DOMAIN_TOKEN));
-	else
-		len = scnprintf(buf, buflen - ret, "context mbox %d ",
-				context);
-	buf += len;
-	ret += len;
-	len = scnprintf(buf, buflen - ret, "vcid %u %s%s\n",
-			group->vcid, group->dev_inaccessible ? "i" : "",
-			group->ext_mailbox ? "x" : "");
-	buf += len;
-	ret += len;
-
-	for_each_list_group_client(lc, group) {
-		len = scnprintf(buf, buflen - ret, "client %s %d:%d\n",
-				lc->client->etiface->name,
-				lc->client->pid, lc->client->tgid);
-		buf += len;
-		ret += len;
-	}
-
-	len = scnprintf(buf, buflen - ret, "mappings %zd %zdB\n",
-			group->host_mappings.count +
-			group->dmabuf_mappings.count,
-			edgetpu_group_mappings_total_size(group));
-	buf += len;
-	ret += len;
-	return ret;
-}
-
-static ssize_t groups_show(
-		struct device *dev, struct device_attribute *attr,
-		char *buf)
-{
-	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
-	struct edgetpu_device_group *group;
-	struct edgetpu_list_group *lg;
-	ssize_t ret = 0;
-
-	mutex_lock(&etdev->groups_lock);
-	etdev_for_each_group(etdev, lg, group) {
-		edgetpu_device_group_get(group);
-		ret += show_group(etdev, group, buf + ret, PAGE_SIZE - ret);
-		edgetpu_device_group_put(group);
-	}
-	mutex_unlock(&etdev->groups_lock);
-	return ret;
-}
-static DEVICE_ATTR_RO(groups);
-
 static struct attribute *edgetpu_dev_attrs[] = {
 	&dev_attr_firmware_crash_count.attr,
 	&dev_attr_watchdog_timeout_count.attr,
 	&dev_attr_clients.attr,
-	&dev_attr_groups.attr,
 	NULL,
 };
 
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 26c5eb1..b492eab 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -84,7 +84,6 @@ struct edgetpu_coherent_mem {
 	dma_addr_t dma_addr;	/* DMA handle for downstream IOMMU, if any */
 	tpu_addr_t tpu_addr;	/* DMA handle for TPU internal IOMMU, if any */
 	u64 host_addr;		/* address mapped on host for debugging */
-	u64 phys_addr;		/* physical address, if available */
 	size_t size;
 #ifdef CONFIG_X86
 	bool is_set_uc;		/* memory has been marked uncached on X86 */
@@ -147,7 +146,7 @@ struct edgetpu_list_device_client {
 	struct edgetpu_client *client;
 };
 
-/* loop through etdev->clients (hold clients_lock prior). */
+/* Macro to loop through etdev->clients (hold clients_lock prior). */
 #define for_each_list_device_client(etdev, c)                                  \
 	list_for_each_entry(c, &etdev->clients, list)
 
diff --git a/drivers/edgetpu/edgetpu-iremap-pool.c b/drivers/edgetpu/edgetpu-iremap-pool.c
index b4b0d2a..91b8fea 100644
--- a/drivers/edgetpu/edgetpu-iremap-pool.c
+++ b/drivers/edgetpu/edgetpu-iremap-pool.c
@@ -10,6 +10,7 @@
 #include <linux/genalloc.h>
 #include <linux/kernel.h>
 #include <linux/mm.h>
+#include <linux/mutex.h>
 #include <linux/slab.h>
 
 #include "edgetpu-internal.h"
@@ -23,6 +24,7 @@ struct edgetpu_mempool {
 	tpu_addr_t base_tpu_addr;
 	phys_addr_t base_phys_addr;
 	size_t granule;
+	struct mutex lock;
 };
 
 int edgetpu_iremap_pool_create(struct edgetpu_dev *etdev, void *base_vaddr,
@@ -42,6 +44,8 @@ int edgetpu_iremap_pool_create(struct edgetpu_dev *etdev, void *base_vaddr,
 	if (!pool)
 		return -ENOMEM;
 
+	mutex_init(&pool->lock);
+
 	pool->gen_pool = gen_pool_create(ilog2(granule), -1);
 	if (!pool->gen_pool) {
 		kfree(pool);
@@ -84,21 +88,21 @@ int edgetpu_iremap_alloc(struct edgetpu_dev *etdev, size_t size,
 
 	if (!etmempool)
 		return edgetpu_alloc_coherent(etdev, size, mem, context_id);
-
+	mutex_lock(&etmempool->lock);
 	size = __ALIGN_KERNEL(size, etmempool->granule);
 	addr = gen_pool_alloc(etmempool->gen_pool, size);
-	if (!addr)
+	if (!addr) {
+		mutex_unlock(&etmempool->lock);
 		return -ENOMEM;
-
+	}
 	mem->vaddr = (void *)addr;
 	offset = mem->vaddr - etmempool->base_vaddr;
 	mem->dma_addr = etmempool->base_dma_addr + offset;
 	mem->tpu_addr = etmempool->base_tpu_addr + offset;
-	mem->phys_addr = etmempool->base_phys_addr + offset;
 	mem->size = size;
 	etdev_dbg(etdev, "%s @ %llx IOVA = %llx size = %zu",
 		   __func__, (u64)mem->vaddr, mem->dma_addr, size);
-
+	mutex_unlock(&etmempool->lock);
 	return 0;
 }
 
@@ -112,12 +116,13 @@ void edgetpu_iremap_free(struct edgetpu_dev *etdev,
 		edgetpu_free_coherent(etdev, mem, context_id);
 		return;
 	}
-
+	mutex_lock(&etmempool->lock);
 	etdev_dbg(etdev, "%s @ %llx IOVA = %llx size = %zu",
 		  __func__, (u64)mem->vaddr, mem->dma_addr, mem->size);
 	gen_pool_free(etmempool->gen_pool, (unsigned long)mem->vaddr,
 		      mem->size);
 	mem->vaddr = NULL;
+	mutex_unlock(&etmempool->lock);
 }
 
 int edgetpu_iremap_mmap(struct edgetpu_dev *etdev, struct vm_area_struct *vma,
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index d32be92..b54fed1 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -951,7 +951,7 @@ void edgetpu_kci_mappings_show(struct edgetpu_dev *etdev, struct seq_file *s)
 	if (!kci || !kci->mailbox)
 		return;
 
-	seq_printf(s, "kci context mbox %u:\n", EDGETPU_CONTEXT_KCI);
+	seq_printf(s, "kci context %u:\n", EDGETPU_CONTEXT_KCI);
 	seq_printf(s, "  0x%llx %lu cmdq - %pad\n",
 		   kci->cmd_queue_mem.tpu_addr,
 		   DIV_ROUND_UP(
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 0785e03..9bda7d6 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -999,7 +999,8 @@ static int edgetpu_mailbox_external_alloc_enable(struct edgetpu_client *client,
 	if (ret) {
 		while (i--) {
 			id = ext_mailbox->descriptors[i].mailbox->mailbox_id;
-			edgetpu_mailbox_deactivate(group->etdev, id);
+			if (edgetpu_mailbox_deactivate(group->etdev, id))
+				etdev_err(group->etdev, "Deactivate mailbox %d failed", id);
 		}
 		/*
 		 * Deactivate only fails if f/w is unresponsive which will put group
@@ -1045,7 +1046,8 @@ void edgetpu_mailbox_external_disable_free_locked(struct edgetpu_device_group *g
 	for (i = 0; i < ext_mailbox->count; i++) {
 		id = ext_mailbox->descriptors[i].mailbox->mailbox_id;
 		etdev_dbg(group->etdev, "Disabling mailbox: %d\n", id);
-		edgetpu_mailbox_deactivate(group->etdev, id);
+		if (edgetpu_mailbox_deactivate(group->etdev, id))
+			etdev_err(group->etdev, "Deactivate mailbox %d failed", id);
 	}
 	/*
 	 * Deactivate only fails if f/w is unresponsive which will put group
@@ -1086,13 +1088,8 @@ out:
 
 int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id)
 {
-	int ret = 0;
+	int ret;
 
-	/*
-	 * A successful enable_ext() increases the wakelock event which prevents wakelock being
-	 * released, so theoretically the check fail here can only happen when enable_ext() is
-	 * failed or not called before.
-	 */
 	if (!edgetpu_wakelock_lock(client->wakelock)) {
 		etdev_err(client->etdev, "Disabling mailbox %d needs wakelock acquired\n",
 			  mailbox_id);
@@ -1106,7 +1103,10 @@ int edgetpu_mailbox_disable_ext(struct edgetpu_client *client, int mailbox_id)
 	}
 	etdev_dbg(client->etdev, "Disabling mailbox: %d\n", mailbox_id);
 
-	edgetpu_mailbox_deactivate(client->etdev, mailbox_id);
+	ret = edgetpu_mailbox_deactivate(client->etdev, mailbox_id);
+	if (ret)
+		etdev_err(client->etdev, "Deactivate mailbox %d failed: %d", mailbox_id, ret);
+
 out:
 	if (!ret)
 		edgetpu_wakelock_dec_event_locked(client->wakelock,
@@ -1139,7 +1139,7 @@ int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, s16 vcid
 	return ret;
 }
 
-void edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id)
+int edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id)
 {
 	struct edgetpu_handshake *eh = &etdev->mailbox_manager->open_devices;
 	const u32 bit = BIT(mailbox_id);
@@ -1148,15 +1148,12 @@ void edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id)
 	mutex_lock(&eh->lock);
 	if (bit & eh->fw_state)
 		ret = edgetpu_kci_close_device(etdev->kci, mailbox_id);
-	if (ret)
-		etdev_err(etdev, "Deactivate mailbox %d failed: %d", mailbox_id, ret);
-	/*
-	 * Always clears the states, FW should never reject CLOSE_DEVICE requests unless it's
-	 * unresponsive.
-	 */
-	eh->state &= ~bit;
-	eh->fw_state &= ~bit;
+	if (!ret) {
+		eh->state &= ~bit;
+		eh->fw_state &= ~bit;
+	}
 	mutex_unlock(&eh->lock);
+	return ret;
 }
 
 void edgetpu_handshake_clear_fw_state(struct edgetpu_handshake *eh)
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index d76534a..3fd961f 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -362,7 +362,7 @@ int edgetpu_mailbox_activate(struct edgetpu_dev *etdev, u32 mailbox_id, s16 vcid
 /*
  * Similar to edgetpu_mailbox_activate() but sends CLOSE_DEVICE KCI instead.
  */
-void edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id);
+int edgetpu_mailbox_deactivate(struct edgetpu_dev *etdev, u32 mailbox_id);
 /* Sets @eh->fw_state to 0. */
 void edgetpu_handshake_clear_fw_state(struct edgetpu_handshake *eh);
 /*
diff --git a/drivers/edgetpu/edgetpu-mapping.c b/drivers/edgetpu/edgetpu-mapping.c
index fbffcce..136cf29 100644
--- a/drivers/edgetpu/edgetpu-mapping.c
+++ b/drivers/edgetpu/edgetpu-mapping.c
@@ -152,21 +152,3 @@ void edgetpu_mappings_show(struct edgetpu_mapping_root *mappings,
 
 	edgetpu_mapping_unlock(mappings);
 }
-
-size_t edgetpu_mappings_total_size(struct edgetpu_mapping_root *mappings)
-{
-	struct rb_node *node;
-	size_t total = 0;
-
-	edgetpu_mapping_lock(mappings);
-
-	for (node = rb_first(&mappings->rb); node; node = rb_next(node)) {
-		struct edgetpu_mapping *map =
-			container_of(node, struct edgetpu_mapping, node);
-
-		total += map->map_size;
-	}
-
-	edgetpu_mapping_unlock(mappings);
-	return total;
-}
diff --git a/drivers/edgetpu/edgetpu-mapping.h b/drivers/edgetpu/edgetpu-mapping.h
index 8f075aa..8ad165b 100644
--- a/drivers/edgetpu/edgetpu-mapping.h
+++ b/drivers/edgetpu/edgetpu-mapping.h
@@ -40,8 +40,6 @@ struct edgetpu_mapping {
 	u64 host_address;
 	u32 die_index; /* this mapping is mapped on the @die_index-th die */
 	tpu_addr_t device_address;
-	/* Size of buffer mapped in bytes. Always set. */
-	size_t map_size;
 	/*
 	 * The size used for allocating @alloc_iova in bytes. This field may be
 	 * set by edgetpu_mmu_map().
@@ -176,7 +174,4 @@ static inline int mmu_flag_to_iommu_prot(u32 mmu_flags, struct device *dev,
 	return prot;
 }
 
-/* Return total size of mappings under the supplied root. */
-size_t edgetpu_mappings_total_size(struct edgetpu_mapping_root *mappings);
-
 #endif /* __EDGETPU_MAPPING_H__ */
diff --git a/drivers/edgetpu/edgetpu-shared-fw.c b/drivers/edgetpu/edgetpu-shared-fw.c
new file mode 100644
index 0000000..c41215e
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-shared-fw.c
@@ -0,0 +1,284 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Edge TPU shared firmware management.
+ *
+ * Copyright (C) 2020 Google, Inc.
+ */
+
+#include <linux/device.h>
+#include <linux/firmware.h>
+#include <linux/list.h>
+#include <linux/mutex.h>
+#include <linux/refcount.h>
+#include <linux/slab.h>
+#include <linux/string.h>
+
+#include "edgetpu-firmware.h"
+#include "edgetpu-firmware-util.h"
+#include "edgetpu-internal.h"
+#include "edgetpu-shared-fw.h"
+
+struct edgetpu_shared_fw_buffer {
+	/*
+	 * Shared firmware buffer is managed by `global.firmware_list`, so that
+	 * each data member is protected by `global.lock`.
+	 */
+	struct list_head list;
+	/*
+	 * Counting for devices holding the buffer. We can only release the data
+	 * buffer if there is no device nor sysfs holding the firmware.
+	 *
+	 * Even when the reference count atomically decreased down to 0, there's
+	 * a chance that someone is traversing list and trying to read this
+	 * `ref`.  So `ref` must still be protected by `glock.lock` in this
+	 * case.
+	 */
+	refcount_t ref;
+	/*
+	 * Indicates if this buffer is loaded by sysfs.
+	 *
+	 * Reference count caused by sysfs load should be exactly 1, and we can
+	 * only unload firmware by sysfs if already loaded by sysfs.
+	 */
+	bool is_sysfs_loaded;
+	/* Firmware name, the same as that loaded by request_firmware() API. */
+	const char *name;
+	void *vaddr;
+	/* The size of buffer is aligned to `global.init_data.size_align`. */
+	size_t size;
+};
+
+const char *
+edgetpu_shared_fw_buffer_name(const struct edgetpu_shared_fw_buffer *buffer)
+{
+	return buffer->name;
+}
+
+void *
+edgetpu_shared_fw_buffer_vaddr(const struct edgetpu_shared_fw_buffer *buffer)
+{
+	return buffer->vaddr;
+}
+
+size_t
+edgetpu_shared_fw_buffer_size(const struct edgetpu_shared_fw_buffer *buffer)
+{
+	return buffer->size;
+}
+
+/*
+ * Lock protected global data.
+ *
+ * global.lock is required for invoking _locked functions in this file.
+ */
+static struct {
+	struct mutex lock;
+	struct edgetpu_shared_fw_init_data init_data;
+	struct list_head firmware_list;
+} global = {
+	.lock = __MUTEX_INITIALIZER(global.lock),
+	.firmware_list = LIST_HEAD_INIT(global.firmware_list),
+};
+
+#define for_each_shared_fw_buffer(buffer) \
+	list_for_each_entry(buffer, &global.firmware_list, list)
+#define for_each_shared_fw_buffer_safe(cur_buf, nxt_buf) \
+	list_for_each_entry_safe(cur_buf, nxt_buf, &global.firmware_list, list)
+
+static struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_find_locked(const char *name)
+{
+	struct edgetpu_shared_fw_buffer *buffer;
+
+	for_each_shared_fw_buffer(buffer) {
+		if (!strcmp(name, buffer->name))
+			return buffer;
+	}
+	return NULL;
+}
+
+void
+edgetpu_shared_fw_init(const struct edgetpu_shared_fw_init_data *init_data)
+{
+	if (list_empty(&global.firmware_list))
+		global.init_data = *init_data;
+}
+
+void edgetpu_shared_fw_exit(void)
+{
+	struct edgetpu_shared_fw_buffer *cur_buf, *nxt_buf;
+
+	mutex_lock(&global.lock);
+
+	if (!list_empty(&global.firmware_list))
+		pr_warn("%s: firmware not released on exiting\n", __func__);
+
+	for_each_shared_fw_buffer_safe(cur_buf, nxt_buf)
+		list_del(&cur_buf->list);
+
+	mutex_unlock(&global.lock);
+}
+
+static struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_get_locked(struct edgetpu_shared_fw_buffer *buffer)
+{
+	if (!buffer)
+		return NULL;
+	if (!refcount_inc_not_zero(&buffer->ref))
+		return NULL;
+	return buffer;
+}
+
+struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_get(struct edgetpu_shared_fw_buffer *buffer)
+{
+	mutex_lock(&global.lock);
+	buffer = edgetpu_shared_fw_get_locked(buffer);
+	mutex_unlock(&global.lock);
+	return buffer;
+}
+
+struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_get_by_name(const char *name)
+{
+	struct edgetpu_shared_fw_buffer *buffer;
+
+	mutex_lock(&global.lock);
+	buffer = edgetpu_shared_fw_get_locked(
+			edgetpu_shared_fw_find_locked(name));
+	mutex_unlock(&global.lock);
+	return buffer;
+}
+
+static struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_load_locked(const char *name, struct edgetpu_dev *etdev)
+{
+	int ret;
+	const struct firmware *fw;
+	size_t aligned_size;
+	struct edgetpu_shared_fw_buffer *buffer;
+
+	buffer = edgetpu_shared_fw_get_locked(
+			edgetpu_shared_fw_find_locked(name));
+	if (buffer) {
+		pr_debug("%s: found shared fw image %s\n", __func__, name);
+		return buffer;
+	}
+
+	pr_debug("%s: shared fw image %s not found, requesting\n",
+		 __func__, name);
+	ret = request_firmware(&fw, name, etdev ? etdev->etiface->etcdev : NULL);
+	if (ret)
+		goto out;
+
+	aligned_size = ALIGN(fw->size, global.init_data.size_align);
+
+	buffer = kzalloc(sizeof(*buffer), GFP_KERNEL);
+	if (!buffer) {
+		ret = -ENOMEM;
+		goto out_release_firmware;
+	}
+
+	buffer->name = kstrdup(name, GFP_KERNEL);
+	if (!buffer->name) {
+		ret = -ENOMEM;
+		goto out_kfree_buffer;
+	}
+
+	/* Allocated in page alignment for mmu and dma mapping. */
+	if (aligned_size < PAGE_SIZE)
+		buffer->vaddr = kmalloc_order(aligned_size, GFP_KERNEL, 1);
+	else
+		buffer->vaddr = kmalloc(aligned_size, GFP_KERNEL);
+	if (!buffer->vaddr) {
+		ret = -ENOMEM;
+		goto out_kfree_buffer_name;
+	}
+	memcpy(buffer->vaddr, fw->data, fw->size);
+	release_firmware(fw);
+
+	buffer->size = aligned_size;
+	refcount_set(&buffer->ref, 1);
+
+	list_add(&buffer->list, &global.firmware_list);
+	return buffer;
+
+out_kfree_buffer_name:
+	kfree(buffer->name);
+out_kfree_buffer:
+	kfree(buffer);
+out_release_firmware:
+	release_firmware(fw);
+out:
+	return ERR_PTR(ret);
+}
+
+struct edgetpu_shared_fw_buffer *edgetpu_shared_fw_load(
+	const char *name, struct edgetpu_dev *etdev)
+{
+	struct edgetpu_shared_fw_buffer *buffer;
+
+	mutex_lock(&global.lock);
+	buffer = edgetpu_shared_fw_load_locked(name, etdev);
+	mutex_unlock(&global.lock);
+	return buffer;
+}
+
+int edgetpu_firmware_chip_load_locked(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_desc *fw_desc, const char *name)
+{
+	int ret;
+	struct edgetpu_dev *etdev = et_fw->etdev;
+	struct edgetpu_shared_fw_buffer *shared_buf;
+
+	shared_buf = edgetpu_shared_fw_load(name, etdev);
+	if (IS_ERR(shared_buf)) {
+		ret = PTR_ERR(shared_buf);
+		etdev_dbg(etdev, "shared buffer loading failed: %d\n", ret);
+		return ret;
+	}
+	fw_desc->shared_buf = shared_buf;
+	fw_desc->buf.vaddr = edgetpu_shared_fw_buffer_vaddr(shared_buf);
+	fw_desc->buf.alloc_size = edgetpu_shared_fw_buffer_size(shared_buf);
+	fw_desc->buf.used_size = fw_desc->buf.alloc_size;
+	fw_desc->buf.name = edgetpu_shared_fw_buffer_name(shared_buf);
+	return 0;
+}
+
+void edgetpu_firmware_chip_unload_locked(
+		struct edgetpu_firmware *et_fw,
+		struct edgetpu_firmware_desc *fw_desc)
+{
+	fw_desc->buf.vaddr = NULL;
+	fw_desc->buf.alloc_size = 0;
+	fw_desc->buf.used_size = 0;
+	fw_desc->buf.name = NULL;
+	edgetpu_shared_fw_put(fw_desc->shared_buf);
+	fw_desc->shared_buf = NULL;
+}
+
+static void
+edgetpu_shared_fw_put_locked(struct edgetpu_shared_fw_buffer *buffer)
+{
+	if (!buffer)
+		return;
+
+	/*
+	 * buffer->ref IS protected by global.lock. See also `ref` in `struct
+	 * edgetpu_shared_fw_buffer`.
+	 */
+	if (refcount_dec_and_test(&buffer->ref)) {
+		kfree(buffer->vaddr);
+		kfree(buffer->name);
+		list_del(&buffer->list);
+		kfree(buffer);
+	}
+}
+
+void edgetpu_shared_fw_put(struct edgetpu_shared_fw_buffer *buffer)
+{
+	mutex_lock(&global.lock);
+	edgetpu_shared_fw_put_locked(buffer);
+	mutex_unlock(&global.lock);
+}
diff --git a/drivers/edgetpu/edgetpu-shared-fw.h b/drivers/edgetpu/edgetpu-shared-fw.h
new file mode 100644
index 0000000..28033e3
--- /dev/null
+++ b/drivers/edgetpu/edgetpu-shared-fw.h
@@ -0,0 +1,74 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Implements utilities for shared firmware management of EdgeTPU.
+ *
+ * Copyright (C) 2020 Google, Inc.
+ */
+#ifndef __EDGETPU_SHARED_FW_H__
+#define __EDGETPU_SHARED_FW_H__
+
+#include <linux/device.h>
+
+#include "edgetpu-internal.h"
+
+struct edgetpu_shared_fw_buffer;
+
+/*
+ * name for this firmware in null terminated string, the same as which loaded by
+ * linux request_firmware API
+ */
+const char *
+edgetpu_shared_fw_buffer_name(const struct edgetpu_shared_fw_buffer *buffer);
+/* host address for this firmware */
+void *
+edgetpu_shared_fw_buffer_vaddr(const struct edgetpu_shared_fw_buffer *buffer);
+/* size in bytes for this firmware */
+size_t
+edgetpu_shared_fw_buffer_size(const struct edgetpu_shared_fw_buffer *buffer);
+
+struct edgetpu_shared_fw_init_data {
+	/* firmware size alignment in bytes */
+	size_t size_align;
+};
+
+/* Initializes structures for shared firmware management. */
+void
+edgetpu_shared_fw_init(const struct edgetpu_shared_fw_init_data *init_data);
+/* Finalizes structures for shared firmware management. */
+void edgetpu_shared_fw_exit(void);
+
+/*
+ * Load reference counted shared firmware from file system. Increase reference
+ * count by 1 if the firmware is already loaded before.
+ *
+ * Firmware loaded by this function should be released by
+ * edgetpu_shared_fw_put().
+ *
+ * @name: firmware path to be loaded
+ * @etdev: requesting edgetpu_dev, if any, for logging
+ */
+struct edgetpu_shared_fw_buffer *edgetpu_shared_fw_load(
+	const char *name, struct edgetpu_dev *etdev);
+
+/*
+ * Increase the reference count of the buffer by 1.
+ *
+ * returns the buffer, behave the same as other *_get/put() functions
+ */
+struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_get(struct edgetpu_shared_fw_buffer *buffer);
+/*
+ * Find the shared firmware by name and increase the reference count of the
+ * found buffer by 1.
+ *
+ * returns NULL on error or not found
+ */
+struct edgetpu_shared_fw_buffer *
+edgetpu_shared_fw_get_by_name(const char *name);
+
+/*
+ * Decrease the reference count by 1 and free the shared buffer if its
+ * reference count reaches 0.
+ */
+void edgetpu_shared_fw_put(struct edgetpu_shared_fw_buffer *buffer);
+#endif /* __EDGETPU_SHARED_FW_H__ */
diff --git a/drivers/edgetpu/edgetpu-thermal.h b/drivers/edgetpu/edgetpu-thermal.h
index dbd283f..8d4b23f 100644
--- a/drivers/edgetpu/edgetpu-thermal.h
+++ b/drivers/edgetpu/edgetpu-thermal.h
@@ -23,7 +23,6 @@ struct edgetpu_thermal {
 	struct mutex lock;
 	void *op_data;
 	unsigned long cooling_state;
-	unsigned long sysfs_req;
 	unsigned int tpu_num_states;
 	struct edgetpu_dev *etdev;
 	bool thermal_suspended; /* TPU thermal suspended state */
diff --git a/drivers/edgetpu/include/linux/platform_data/sscoredump.h b/drivers/edgetpu/include/linux/platform_data/sscoredump.h
new file mode 100644
index 0000000..fee2872
--- /dev/null
+++ b/drivers/edgetpu/include/linux/platform_data/sscoredump.h
@@ -0,0 +1,35 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+#ifndef __SUBSYSTEM_COREDUMP_H
+#define __SUBSYSTEM_COREDUMP_H
+
+#include <linux/device.h>
+#include <linux/module.h>
+#include <linux/vmalloc.h>
+#include <linux/platform_device.h>
+
+#define SSCD_NAME			"sscoredump"
+
+/* sscd segment ( similar to ELF memory segments) */
+struct sscd_segment {
+	void		*addr;
+	u64		size;
+	u64		flags;
+
+	/* passed to elf sprogram header */
+	void		*paddr;
+	void		*vaddr;
+} __packed;
+
+/* sscd_report flags */
+#define SSCD_FLAGS_ELFARM32HDR		0x0001
+#define SSCD_FLAGS_ELFARM64HDR		0x0002
+
+
+struct sscd_platform_data {
+	/* report crash */
+	int (*sscd_report)(struct platform_device *pdev,
+			   struct sscd_segment *segs, int nsegs,
+			   u64 flags, const char *crash_info);
+};
+
+#endif /* __SUBSYSTEM_COREDUMP_H */
diff --git a/drivers/edgetpu/janeiro-device.c b/drivers/edgetpu/janeiro-device.c
index 36dbe76..7b1aa6d 100644
--- a/drivers/edgetpu/janeiro-device.c
+++ b/drivers/edgetpu/janeiro-device.c
@@ -15,14 +15,6 @@
 #include "janeiro-platform.h"
 #include "mobile-pm.h"
 
-#define SSMT_NS_READ_STREAM_VID_OFFSET(n) (0x1000u + (0x4u * (n)))
-#define SSMT_NS_WRITE_STREAM_VID_OFFSET(n) (0x1200u + (0x4u * (n)))
-
-#define SSMT_NS_READ_STREAM_VID_REG(base, n)                                   \
-	((base) + SSMT_NS_READ_STREAM_VID_OFFSET(n))
-#define SSMT_NS_WRITE_STREAM_VID_REG(base, n)                                  \
-	((base) + SSMT_NS_WRITE_STREAM_VID_OFFSET(n))
-
 static irqreturn_t janeiro_mailbox_handle_irq(struct edgetpu_dev *etdev,
 					      int irq)
 {
@@ -59,6 +51,10 @@ irqreturn_t edgetpu_chip_irq_handler(int irq, void *arg)
 	struct edgetpu_dev *etdev = arg;
 
 	edgetpu_telemetry_irq_handler(etdev);
+	/*
+	 * use this as HOST_NONSECURE_INT_SRC_STATUS_REG not present in
+	 * Janeiro.
+	 */
 	return janeiro_mailbox_handle_irq(etdev, irq);
 }
 
@@ -69,17 +65,6 @@ u64 edgetpu_chip_tpu_timestamp(struct edgetpu_dev *etdev)
 
 void edgetpu_chip_init(struct edgetpu_dev *etdev)
 {
-	int i;
-	struct janeiro_platform_dev *jpdev = to_janeiro_dev(etdev);
-
-	if (!jpdev->ssmt_base)
-		return;
-
-	/* Setup non-secure SCIDs, assume VID = SCID */
-	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
-		writel(i, SSMT_NS_READ_STREAM_VID_REG(jpdev->ssmt_base, i));
-		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(jpdev->ssmt_base, i));
-	}
 }
 
 void edgetpu_chip_exit(struct edgetpu_dev *etdev)
diff --git a/drivers/edgetpu/janeiro-firmware.c b/drivers/edgetpu/janeiro-firmware.c
index cbc7cad..3ce8880 100644
--- a/drivers/edgetpu/janeiro-firmware.c
+++ b/drivers/edgetpu/janeiro-firmware.c
@@ -133,7 +133,7 @@ static int janeiro_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 {
 	int ret = 0;
 	void *image_vaddr;
-	u32 tpu_addr, phys_addr, size, i, j;
+	u32 tpu_addr, phys_addr, size, i;
 	struct janeiro_image_config *image_config;
 	struct janeiro_firmware_data *data;
 	struct edgetpu_dev *etdev = et_fw->etdev;
@@ -168,10 +168,8 @@ static int janeiro_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 
 		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
 	}
-	for (i = 0, j = 0; i < MAX_IOMMU_MAPPINGS; i++) {
+	for (i = 0; i < image_config->num_iommu_mapping; i++) {
 		tpu_addr = image_config->mappings[i].virt_address;
-		if (!tpu_addr)
-			continue;
 		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
 		phys_addr = (image_config->mappings[i].image_config_value & ~(0xFFF));
 
@@ -183,17 +181,11 @@ static int janeiro_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 				  ret, tpu_addr, phys_addr, size);
 			goto err;
 		}
-		data->mappings[j].virt_address = tpu_addr;
-		data->mappings[j++].image_config_value =
-						image_config->mappings[i].image_config_value;
+		data->mappings[i].virt_address = tpu_addr;
+		data->mappings[i].image_config_value = image_config->mappings[i].image_config_value;
 	}
 
-	if (image_config->num_iommu_mapping != j) {
-		etdev_err(etdev, "Invalid firmware header\n");
-		ret = -EINVAL;
-		goto err;
-	}
-	data->num_mapping = j;
+	data->num_mapping = image_config->num_iommu_mapping;
 
 	/* Skip the header */
 	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
@@ -201,9 +193,9 @@ static int janeiro_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 	memunmap(image_vaddr);
 	return 0;
 err:
-	while (j--) {
-		tpu_addr = data->mappings[j].virt_address;
-		size = CONFIG_TO_SIZE(data->mappings[j].image_config_value);
+	while (i--) {
+		tpu_addr = data->mappings[i].virt_address;
+		size = CONFIG_TO_SIZE(data->mappings[i].image_config_value);
 		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
 	}
 	data->num_mapping = 0;
diff --git a/drivers/edgetpu/janeiro-platform.c b/drivers/edgetpu/janeiro-platform.c
index 9507e4a..73c6119 100644
--- a/drivers/edgetpu/janeiro-platform.c
+++ b/drivers/edgetpu/janeiro-platform.c
@@ -117,29 +117,6 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
-static int janeiro_parse_ssmt(struct janeiro_platform_dev *etpdev)
-{
-	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
-	struct platform_device *pdev = to_platform_device(etdev->dev);
-	struct resource *res;
-	int rc;
-	void __iomem *ssmt_base;
-
-	res = platform_get_resource_byname(pdev, IORESOURCE_MEM, "ssmt");
-	if (!res) {
-		etdev_warn(etdev, "Failed to find SSMT register base");
-		return -EINVAL;
-	}
-	ssmt_base = devm_ioremap_resource(&pdev->dev, res);
-	if (IS_ERR(ssmt_base)) {
-		rc = PTR_ERR(ssmt_base);
-		etdev_warn(etdev, "Failed to map SSMT register base: %d\n", rc);
-		return rc;
-	}
-	etpdev->ssmt_base = ssmt_base;
-	return 0;
-}
-
 /*
  * Set shareability for enabling IO coherency in Janeiro
  */
@@ -286,13 +263,6 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 		goto out_remove_device;
 	}
 
-	ret = janeiro_parse_ssmt(edgetpu_pdev);
-	if (ret)
-		dev_warn(
-			dev,
-			"SSMT setup failed (%d). Context isolation not enforced\n",
-			ret);
-
 	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_LOG,
 				  &edgetpu_pdev->log_mem);
 	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_TRACE,
diff --git a/drivers/edgetpu/janeiro-platform.h b/drivers/edgetpu/janeiro-platform.h
index 51eb8d6..09991cf 100644
--- a/drivers/edgetpu/janeiro-platform.h
+++ b/drivers/edgetpu/janeiro-platform.h
@@ -11,10 +11,7 @@
 #include <linux/io.h>
 #include <linux/mutex.h>
 #include <linux/types.h>
-
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
 #include <soc/google/bcl.h>
-#endif
 
 #include "edgetpu-internal.h"
 
diff --git a/drivers/edgetpu/janeiro-pm.c b/drivers/edgetpu/janeiro-pm.c
index 700601e..e69f774 100644
--- a/drivers/edgetpu/janeiro-pm.c
+++ b/drivers/edgetpu/janeiro-pm.c
@@ -10,11 +10,8 @@
 #include <linux/iopoll.h>
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
-#include <linux/version.h>
-
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
 #include <soc/google/bcl.h>
-#endif
+#include <linux/version.h>
 
 #include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
@@ -28,9 +25,9 @@
 
 #define SHUTDOWN_DELAY_US_MIN		20
 #define SHUTDOWN_DELAY_US_MAX		20
-#define BOOTUP_DELAY_US_MIN		200
-#define BOOTUP_DELAY_US_MAX		250
-#define SHUTDOWN_MAX_DELAY_COUNT	50
+#define BOOTUP_DELAY_US_MIN		100
+#define BOOTUP_DELAY_US_MAX		150
+#define SHUTDOWN_MAX_DELAY_COUNT	20
 
 /* Default power state */
 static int power_state = TPU_ACTIVE_NOM;
@@ -115,8 +112,7 @@ static int janeiro_pwr_state_set_locked(void *data, u64 val)
 		do {
 			/* Delay 20us per retry till blk shutdown finished */
 			usleep_range(SHUTDOWN_DELAY_US_MIN, SHUTDOWN_DELAY_US_MAX);
-			/* Only poll for BLK status instead of CLK rate */
-			curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 1);
+			curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
 			if (!curr_state)
 				break;
 			timeout_cnt++;
@@ -241,14 +237,14 @@ static int janeiro_set_lpm(struct edgetpu_dev *etdev)
 
 	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_START, 1);
 	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM0_STATUS, val,
-				 val & 0x80, 20, EDGETPU_LPM_CHANGE_TIMEOUT);
+				 val & 0x80, 5, EDGETPU_LPM_CHANGE_TIMEOUT);
 	if (ret) {
 		etdev_err(etdev, "Set LPM0 failed: %d\n", ret);
 		return ret;
 	}
 	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM1_START, 1);
 	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM1_STATUS, val,
-				 val & 0x80, 20, EDGETPU_LPM_CHANGE_TIMEOUT);
+				 val & 0x80, 5, EDGETPU_LPM_CHANGE_TIMEOUT);
 	if (ret) {
 		etdev_err(etdev, "Set LPM1 failed: %d\n", ret);
 		return ret;
diff --git a/drivers/edgetpu/janeiro-pm.h b/drivers/edgetpu/janeiro-pm.h
new file mode 100644
index 0000000..78991cc
--- /dev/null
+++ b/drivers/edgetpu/janeiro-pm.h
@@ -0,0 +1,79 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Power management header for Janeiro.
+ *
+ * Copyright (C) 2020 Google, Inc.
+ */
+#ifndef __JANEIRO_PM_H__
+#define __JANEIRO_PM_H__
+
+#include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
+
+/* Can't build out of tree with acpm_dvfs unless kernel supports ACPM */
+#if IS_ENABLED(CONFIG_ACPM_DVFS)
+
+#include <linux/acpm_dvfs.h>
+
+#else
+
+static unsigned long exynos_acpm_rate;
+static inline int exynos_acpm_set_rate(unsigned int id, unsigned long rate)
+{
+	exynos_acpm_rate = rate;
+	return 0;
+}
+static inline int exynos_acpm_set_init_freq(unsigned int dfs_id,
+					    unsigned long freq)
+{
+	return 0;
+}
+static inline unsigned long exynos_acpm_get_rate(unsigned int id,
+						 unsigned long dbg_val)
+{
+	return exynos_acpm_rate;
+}
+static inline int exynos_acpm_set_policy(unsigned int id, unsigned long policy)
+{
+	return 0;
+}
+#endif /* IS_ENABLED(CONFIG_ACPM_DVFS) */
+//TODO(b/185797093): check abrolhos ported values for janeiro
+/*
+ * TPU Power States:
+ * 0:		Off
+ * 227000	Ultra Underdrive @227MHz
+ * 625000:	Super Underdrive @625MHz
+ * 845000:	Underdrive @845MHz
+ * 1066000:	Nominal @1066MHz
+ */
+enum tpu_pwr_state {
+	TPU_OFF = 0,
+	TPU_ACTIVE_UUD = 227000,
+	TPU_ACTIVE_SUD = 625000,
+	TPU_ACTIVE_UD  = 845000,
+	TPU_ACTIVE_NOM = 1066000,
+};
+
+/*
+ * Request codes from firmware
+ * Values must match with firmware code base
+ */
+enum janeiro_reverse_kci_code {
+	RKCI_CODE_PM_QOS = RKCI_CHIP_CODE_FIRST + 1,
+	RKCI_CODE_BTS = RKCI_CHIP_CODE_FIRST + 2,
+};
+
+#define TPU_POLICY_MAX	TPU_ACTIVE_NOM
+
+#define TPU_ACPM_DOMAIN			7
+
+int janeiro_pm_create(struct edgetpu_dev *etdev);
+
+void janeiro_pm_destroy(struct edgetpu_dev *etdev);
+
+void janeiro_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
+
+void janeiro_pm_set_bts(struct edgetpu_dev *etdev, u32 bts_val);
+
+#endif /* __JANEIRO_PM_H__ */
diff --git a/drivers/edgetpu/janeiro/config-mailbox.h b/drivers/edgetpu/janeiro/config-mailbox.h
index 6d7bce9..e2a07cf 100644
--- a/drivers/edgetpu/janeiro/config-mailbox.h
+++ b/drivers/edgetpu/janeiro/config-mailbox.h
@@ -30,6 +30,10 @@
 #define JANEIRO_CSR_MBOX_CMD_QUEUE_DOORBELL_SET_OFFSET 0x1000
 #define JANEIRO_CSR_MBOX_RESP_QUEUE_DOORBELL_SET_OFFSET 0x1800
 #define EDGETPU_MBOX_BASE JANEIRO_CSR_MBOX2_CONTEXT_ENABLE
+// TODO: check correct values
+/* CSR storing mailbox response queue doorbell status */
+#define HOST_NONSECURE_INT_SRC_STATUS_REG 0x000f0000
+#define HOST_NONSECURE_INT_SRC_CLEAR_REG 0x000f0008
 
 static inline u32 edgetpu_mailbox_get_context_csr_base(u32 index)
 {