Merge branch 'whitechapel' into android-gs-pixel-5.10-gs101

  edgetpu: add iterator for each device in an MCP
  Bug: 174552882
  edgetpu: mobile: fixups for android review comments
  Bug: 205749891
  edgetpu: mcp synchronous usage stats update at device group disband
  Bug: 195612161
  edgetpu: unregister KCI irq handler
  Bug: 206717576
  Bug: 204847583
  edgetpu: add h/w preempt metric
  Bug: 205901661
  edgetpu: define and use vma offset for external test
  Bug: 204846795
  edgetpu: fix compile failure on older linux for pin_user_pages
  Bug: 205000780
  edgetpu: add ioctl for testing interop interface
  Bug: 204846795
  edgetpu: do SSMT setup after firmware image loaded
  Bug: 204155811
  edgetpu: hermosa / shared firmware auth changes from prodkernel
  Bug: 195776283

GitOrigin-RevId: 835d991b33de161b6ce677cd1d190ee097daac7f
Change-Id: I1c7e35550c7c14973358fb86fd26bc4b314a780c

16 files changed, 297 insertions, 83 deletions

diff --git a/drivers/edgetpu/abrolhos-device.c b/drivers/edgetpu/abrolhos-device.c
index c2f118f..df23c42 100644
--- a/drivers/edgetpu/abrolhos-device.c
+++ b/drivers/edgetpu/abrolhos-device.c
@@ -18,14 +18,6 @@
 
 #define HOST_NONSECURE_INTRSRCMASKREG	0x000f0004
 
-#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))
-
 /*
  * The interrupt handler for mailboxes.
  *
@@ -85,20 +77,8 @@ u64 edgetpu_chip_tpu_timestamp(struct edgetpu_dev *etdev)
 
 void edgetpu_chip_init(struct edgetpu_dev *etdev)
 {
-	int i;
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-
 	/* Disable the CustomBlock Interrupt. */
 	edgetpu_dev_write_32(etdev, HOST_NONSECURE_INTRSRCMASKREG, 0x1);
-
-	if (!etmdev->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(etmdev->ssmt_base, i));
-		writel(i, SSMT_NS_WRITE_STREAM_VID_REG(etmdev->ssmt_base, i));
-	}
 }
 
 void edgetpu_chip_exit(struct edgetpu_dev *etdev)
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index f4b89ae..9bb84ac 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -50,6 +50,9 @@ enum edgetpu_vma_type {
 	VMA_VII_CSR,
 	VMA_VII_CMDQ,
 	VMA_VII_RESPQ,
+	VMA_EXT_CSR,
+	VMA_EXT_CMDQ,
+	VMA_EXT_RESPQ,
 	/* For VMA_LOG and VMA_TRACE, core id is stored in bits higher than VMA_TYPE_WIDTH. */
 	VMA_LOG,
 	VMA_TRACE,
@@ -103,6 +106,14 @@ static edgetpu_vma_flags_t mmap_vma_flag(unsigned long pgoff)
 		return VMA_VII_CMDQ;
 	case EDGETPU_MMAP_RESP_QUEUE_OFFSET:
 		return VMA_VII_RESPQ;
+#ifdef EDGETPU_FEATURE_INTEROP
+	case EDGETPU_MMAP_EXT_CSR_OFFSET:
+		return VMA_EXT_CSR;
+	case EDGETPU_MMAP_EXT_CMD_QUEUE_OFFSET:
+		return VMA_EXT_CMDQ;
+	case EDGETPU_MMAP_EXT_RESP_QUEUE_OFFSET:
+		return VMA_EXT_RESPQ;
+#endif /* EDGETPU_FEATURE_INTEROP */
 	case EDGETPU_MMAP_LOG_BUFFER_OFFSET:
 		return VMA_DATA_SET(VMA_LOG, 0);
 	case EDGETPU_MMAP_TRACE_BUFFER_OFFSET:
@@ -134,6 +145,12 @@ vma_type_to_wakelock_event(enum edgetpu_vma_type type)
 		return EDGETPU_WAKELOCK_EVENT_CMD_QUEUE;
 	case VMA_VII_RESPQ:
 		return EDGETPU_WAKELOCK_EVENT_RESP_QUEUE;
+	case VMA_EXT_CSR:
+		return EDGETPU_WAKELOCK_EVENT_MBOX_CSR;
+	case VMA_EXT_CMDQ:
+		return EDGETPU_WAKELOCK_EVENT_CMD_QUEUE;
+	case VMA_EXT_RESPQ:
+		return EDGETPU_WAKELOCK_EVENT_RESP_QUEUE;
 	default:
 		return EDGETPU_WAKELOCK_EVENT_END;
 	}
@@ -313,14 +330,22 @@ int edgetpu_mmap(struct edgetpu_client *client, struct vm_area_struct *vma)
 	}
 	switch (type) {
 	case VMA_VII_CSR:
-		ret = edgetpu_mmap_csr(client->group, vma);
+		ret = edgetpu_mmap_csr(client->group, vma, false);
 		break;
 	case VMA_VII_CMDQ:
-		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma);
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma, false);
 		break;
 	case VMA_VII_RESPQ:
-		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE,
-					 vma);
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE, vma, false);
+		break;
+	case VMA_EXT_CSR:
+		ret = edgetpu_mmap_csr(client->group, vma, true);
+		break;
+	case VMA_EXT_CMDQ:
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_CMD_QUEUE, vma, true);
+		break;
+	case VMA_EXT_RESPQ:
+		ret = edgetpu_mmap_queue(client->group, MAILBOX_RESP_QUEUE, vma, true);
 		break;
 	default: /* to appease compiler */
 		break;
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 6f338be..c5d768c 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -7,6 +7,7 @@
 
 #include <linux/atomic.h>
 #include <linux/bitops.h>
+#include <linux/cred.h>
 #include <linux/dma-direction.h>
 #include <linux/dma-mapping.h>
 #include <linux/eventfd.h>
@@ -20,6 +21,7 @@
 #include <linux/seq_file.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>
+#include <linux/uidgid.h>
 
 #include "edgetpu-async.h"
 #include "edgetpu-config.h"
@@ -102,7 +104,7 @@ static int edgetpu_kci_leave_group_worker(struct kci_worker_param *param)
 	struct edgetpu_dev *etdev = edgetpu_device_group_nth_etdev(group, i);
 
 	etdev_dbg(etdev, "%s: leave group %u", __func__, group->workload_id);
-	edgetpu_kci_update_usage_async(etdev);
+	edgetpu_kci_update_usage(etdev);
 	edgetpu_kci_leave_group(etdev->kci);
 	return 0;
 }
@@ -1672,21 +1674,33 @@ void edgetpu_group_mappings_show(struct edgetpu_device_group *group,
 }
 
 int edgetpu_mmap_csr(struct edgetpu_device_group *group,
-		     struct vm_area_struct *vma)
+		     struct vm_area_struct *vma, bool is_external)
 {
 	struct edgetpu_dev *etdev = group->etdev;
 	int ret = 0;
 	ulong phys_base, vma_size, map_size;
 
+	if (is_external && !uid_eq(current_euid(), GLOBAL_ROOT_UID))
+		return -EPERM;
+
 	mutex_lock(&group->lock);
 	if (!edgetpu_group_finalized_and_attached(group)) {
 		ret = edgetpu_group_errno(group);
 		goto out;
 	}
 
+	if (is_external && (!group->ext_mailbox || !group->ext_mailbox->descriptors)) {
+		ret = -ENOENT;
+		goto out;
+	}
+
 	vma_size = vma->vm_end - vma->vm_start;
 	map_size = min(vma_size, USERSPACE_CSR_SIZE);
-	phys_base = etdev->regs.phys + group->vii.mailbox->cmd_queue_csr_base;
+	if (is_external)
+		phys_base = etdev->regs.phys +
+			    group->ext_mailbox->descriptors[0].mailbox->cmd_queue_csr_base;
+	else
+		phys_base = etdev->regs.phys + group->vii.mailbox->cmd_queue_csr_base;
 	ret = io_remap_pfn_range(vma, vma->vm_start, phys_base >> PAGE_SHIFT,
 				 map_size, vma->vm_page_prot);
 	if (ret)
@@ -1699,22 +1713,37 @@ out:
 
 int edgetpu_mmap_queue(struct edgetpu_device_group *group,
 		       enum mailbox_queue_type type,
-		       struct vm_area_struct *vma)
+		       struct vm_area_struct *vma, bool is_external)
 {
 	struct edgetpu_dev *etdev = group->etdev;
 	int ret = 0;
 	edgetpu_queue_mem *queue_mem;
 
+	if (is_external && !uid_eq(current_euid(), GLOBAL_ROOT_UID))
+		return -EPERM;
+
 	mutex_lock(&group->lock);
 	if (!edgetpu_group_finalized_and_attached(group)) {
 		ret = edgetpu_group_errno(group);
 		goto out;
 	}
 
-	if (type == MAILBOX_CMD_QUEUE)
-		queue_mem = &(group->vii.cmd_queue_mem);
-	else
-		queue_mem = &(group->vii.resp_queue_mem);
+	if (is_external && (!group->ext_mailbox || !group->ext_mailbox->descriptors)) {
+		ret = -ENOENT;
+		goto out;
+	}
+
+	if (type == MAILBOX_CMD_QUEUE) {
+		if (is_external)
+			queue_mem = &(group->ext_mailbox->descriptors[0].cmd_queue_mem);
+		else
+			queue_mem = &(group->vii.cmd_queue_mem);
+	} else {
+		if (is_external)
+			queue_mem = &(group->ext_mailbox->descriptors[0].resp_queue_mem);
+		else
+			queue_mem = &(group->vii.resp_queue_mem);
+	}
 
 	if (!queue_mem->vaddr) {
 		ret = -ENXIO;
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 5c5eaaf..bf2776b 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -360,7 +360,7 @@ void edgetpu_group_mappings_show(struct edgetpu_device_group *group,
  * Returns 0 on success.
  */
 int edgetpu_mmap_csr(struct edgetpu_device_group *group,
-		     struct vm_area_struct *vma);
+		     struct vm_area_struct *vma, bool is_external);
 /*
  * Maps the cmd/resp queue memory.
  *
@@ -368,7 +368,7 @@ int edgetpu_mmap_csr(struct edgetpu_device_group *group,
  */
 int edgetpu_mmap_queue(struct edgetpu_device_group *group,
 		       enum mailbox_queue_type type,
-		       struct vm_area_struct *vma);
+		       struct vm_area_struct *vma, bool is_external);
 
 /* Set group eventfd for event notification */
 int edgetpu_group_set_eventfd(struct edgetpu_device_group *group, uint event_id,
diff --git a/drivers/edgetpu/edgetpu-external.c b/drivers/edgetpu/edgetpu-external.c
index 1552774..4b86e13 100644
--- a/drivers/edgetpu/edgetpu-external.c
+++ b/drivers/edgetpu/edgetpu-external.c
@@ -101,10 +101,8 @@ static int edgetpu_external_mailbox_alloc(struct device *edgetpu_dev,
 	}
 
 	ret = edgetpu_mailbox_enable_ext(client, EDGETPU_MAILBOX_ID_USE_ASSOC, &req);
-	if (ret) {
-		edgetpu_device_group_put(group);
-		goto out;
-	}
+	if (ret)
+		goto error_put_group;
 	mutex_lock(&group->lock);
 	ext_mailbox = group->ext_mailbox;
 	if (!ext_mailbox) {
@@ -114,6 +112,7 @@ static int edgetpu_external_mailbox_alloc(struct device *edgetpu_dev,
 	ret = edgetpu_external_mailbox_info_get(info, ext_mailbox);
 unlock:
 	mutex_unlock(&group->lock);
+error_put_group:
 	edgetpu_device_group_put(group);
 out:
 	fdput(f);
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index f24135b..1c7c3c4 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -298,7 +298,7 @@ uint64_t edgetpu_firmware_get_build_time(struct edgetpu_firmware *et_fw);
  */
 bool edgetpu_firmware_verify_signature(struct edgetpu_dev *etdev,
 				       const char *name,
-				       void **image_data, size_t *image_size);
+				       const 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 46de18f..599a6a7 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -6,7 +6,9 @@
  */
 
 #include <linux/atomic.h>
+#include <linux/bitops.h>
 #include <linux/cdev.h>
+#include <linux/cred.h>
 #include <linux/debugfs.h>
 #include <linux/device.h>
 #include <linux/file.h>
@@ -25,6 +27,7 @@
 #include <linux/slab.h>
 #include <linux/types.h>
 #include <linux/uaccess.h>
+#include <linux/uidgid.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-device-group.h"
@@ -39,6 +42,10 @@
 #include "edgetpu-wakelock.h"
 #include "edgetpu.h"
 
+#ifdef EDGETPU_FEATURE_INTEROP
+#include <soc/google/tpu-ext.h>
+#endif
+
 #define CREATE_TRACE_POINTS
 #include <trace/events/edgetpu.h>
 
@@ -685,6 +692,40 @@ static int edgetpu_ioctl_get_fatal_errors(struct edgetpu_client *client,
 	return ret;
 }
 
+#ifdef EDGETPU_FEATURE_INTEROP
+static int edgetpu_ioctl_test_external(struct edgetpu_client *client,
+				       struct edgetpu_test_ext_ioctl __user *argp)
+{
+	int ret = 0;
+	struct edgetpu_test_ext_ioctl test_ext;
+	struct edgetpu_ext_client_info client_info;
+	struct edgetpu_ext_mailbox_info *info;
+
+	if (!uid_eq(current_euid(), GLOBAL_ROOT_UID))
+		return -EPERM;
+
+	if (copy_from_user(&test_ext, argp, sizeof(test_ext)))
+		return -EFAULT;
+
+	if (hweight32(test_ext.mbox_bmap) > 1)
+		return -EINVAL;
+
+	client_info.attr = (struct edgetpu_mailbox_attr __user *)test_ext.attrs;
+	client_info.tpu_fd = test_ext.fd;
+	client_info.mbox_map = test_ext.mbox_bmap;
+
+	info = kmalloc(sizeof(*info) + sizeof(struct edgetpu_ext_mailbox_descriptor *), GFP_KERNEL);
+	if (!info)
+		return -ENOMEM;
+
+	ret = edgetpu_ext_driver_cmd(client->etdev->dev, test_ext.client_type, test_ext.cmd,
+				     &client_info, info);
+
+	kfree(info);
+	return ret;
+}
+#endif /* EDGETPU_FEATURE_INTEROP */
+
 long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 {
 	struct edgetpu_client *client = file->private_data;
@@ -776,7 +817,11 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 	case EDGETPU_GET_FATAL_ERRORS:
 		ret = edgetpu_ioctl_get_fatal_errors(client, argp);
 		break;
-
+#ifdef EDGETPU_FEATURE_INTEROP
+	case EDGETPU_TEST_EXTERNAL:
+		ret = edgetpu_ioctl_test_external(client, argp);
+		break;
+#endif /* EDGETPU_FEATURE_INTEROP */
 	default:
 		return -ENOTTY; /* unknown command */
 	}
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index bfa356f..9339953 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -507,6 +507,8 @@ int edgetpu_kci_init(struct edgetpu_mailbox_manager *mgr,
 int edgetpu_kci_reinit(struct edgetpu_kci *kci)
 {
 	struct edgetpu_mailbox *mailbox = kci->mailbox;
+	struct edgetpu_mailbox_manager *mgr;
+	unsigned long flags;
 	int ret;
 
 	if (!mailbox)
@@ -521,6 +523,13 @@ int edgetpu_kci_reinit(struct edgetpu_kci *kci)
 					QUEUE_SIZE);
 	if (ret)
 		return ret;
+
+	mgr = mailbox->etdev->mailbox_manager;
+	/* Restore KCI irq handler */
+	write_lock_irqsave(&mgr->mailboxes_lock, flags);
+	mailbox->handle_irq = edgetpu_kci_handle_irq;
+	write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
+
 	edgetpu_mailbox_init_doorbells(mailbox);
 	edgetpu_mailbox_enable(mailbox);
 
@@ -529,6 +538,16 @@ int edgetpu_kci_reinit(struct edgetpu_kci *kci)
 
 void edgetpu_kci_cancel_work_queues(struct edgetpu_kci *kci)
 {
+	struct edgetpu_mailbox_manager *mgr;
+	unsigned long flags;
+
+	if (kci->mailbox) {
+		mgr = kci->mailbox->etdev->mailbox_manager;
+		/* Remove IRQ handler to stop responding to interrupts */
+		write_lock_irqsave(&mgr->mailboxes_lock, flags);
+		kci->mailbox->handle_irq = NULL;
+		write_unlock_irqrestore(&mgr->mailboxes_lock, flags);
+	}
 	/* Cancel workers that may send KCIs. */
 	cancel_work_sync(&kci->usage_work);
 	/* Cancel KCI and reverse KCI workers. */
diff --git a/drivers/edgetpu/edgetpu-mcp.h b/drivers/edgetpu/edgetpu-mcp.h
index 1b76ecf..4530d79 100644
--- a/drivers/edgetpu/edgetpu-mcp.h
+++ b/drivers/edgetpu/edgetpu-mcp.h
@@ -99,6 +99,15 @@ int edgetpu_mcp_each(bool stop_on_err, void *data,
 		     int (*callback)(struct edgetpu_mcp *, void *));
 
 /*
+ * Invoke @callback on each device for the specified @mcp.
+ *
+ * @data can be any value, it will be directly passed to @callback.
+ */
+void edgetpu_mcp_foreach_etdev(struct edgetpu_mcp *mcp,
+			       void (*callback)(struct edgetpu_dev *, void *),
+			       void *data);
+
+/*
  * Returns the MCP object @etdev belongs to.
  *
  * Returns NULL when such object is not found.
diff --git a/drivers/edgetpu/edgetpu-usage-stats.c b/drivers/edgetpu/edgetpu-usage-stats.c
index f796117..4c6dfcc 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.c
+++ b/drivers/edgetpu/edgetpu-usage-stats.c
@@ -549,6 +549,27 @@ static ssize_t context_preempt_count_store(struct device *dev,
 static DEVICE_ATTR(context_preempt_count, 0664, context_preempt_count_show,
 		   context_preempt_count_store);
 
+static ssize_t hardware_preempt_count_show(struct device *dev, struct device_attribute *attr,
+					   char *buf)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+	int64_t val;
+
+	val = edgetpu_usage_get_counter(etdev, EDGETPU_COUNTER_HARDWARE_PREEMPTS);
+	return scnprintf(buf, PAGE_SIZE, "%llu\n", val);
+}
+
+static ssize_t hardware_preempt_count_store(struct device *dev, struct device_attribute *attr,
+					    const char *buf, size_t count)
+{
+	struct edgetpu_dev *etdev = dev_get_drvdata(dev);
+
+	edgetpu_counter_clear(etdev, EDGETPU_COUNTER_HARDWARE_PREEMPTS);
+	return count;
+}
+static DEVICE_ATTR(hardware_preempt_count, 0664, hardware_preempt_count_show,
+		   hardware_preempt_count_store);
+
 static ssize_t outstanding_commands_max_show(
 	struct device *dev, struct device_attribute *attr, char *buf)
 {
@@ -659,6 +680,7 @@ static struct attribute *usage_stats_dev_attrs[] = {
 	&dev_attr_param_cache_hit_count.attr,
 	&dev_attr_param_cache_miss_count.attr,
 	&dev_attr_context_preempt_count.attr,
+	&dev_attr_hardware_preempt_count.attr,
 	&dev_attr_outstanding_commands_max.attr,
 	&dev_attr_preempt_depth_max.attr,
 	&dev_attr_fw_thread_stats.attr,
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 20d5ad7..c76899a 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -72,8 +72,10 @@ enum edgetpu_usage_counter_type {
 	EDGETPU_COUNTER_PARAM_CACHE_MISSES = 5,
 	/* Number of times a context got preempted by another. */
 	EDGETPU_COUNTER_CONTEXT_PREEMPTS = 6,
+	/* Number of times a hardware preemption occurred. */
+	EDGETPU_COUNTER_HARDWARE_PREEMPTS = 7,
 
-	EDGETPU_COUNTER_COUNT = 7, /* number of counters above */
+	EDGETPU_COUNTER_COUNT = 8, /* number of counters above */
 };
 
 /* Generic counter. Only reported if it has a value larger than 0. */
diff --git a/drivers/edgetpu/edgetpu.h b/drivers/edgetpu/edgetpu.h
index 6a210fb..ad9a2ce 100644
--- a/drivers/edgetpu/edgetpu.h
+++ b/drivers/edgetpu/edgetpu.h
@@ -11,6 +11,9 @@
 #include <linux/types.h>
 
 /* mmap offsets for mailbox CSRs, command queue, and response queue */
+#define EDGETPU_MMAP_EXT_CSR_OFFSET 0x1500000
+#define EDGETPU_MMAP_EXT_CMD_QUEUE_OFFSET 0x1600000
+#define EDGETPU_MMAP_EXT_RESP_QUEUE_OFFSET 0x1700000
 #define EDGETPU_MMAP_CSR_OFFSET 0x1800000
 #define EDGETPU_MMAP_CMD_QUEUE_OFFSET 0x1900000
 #define EDGETPU_MMAP_RESP_QUEUE_OFFSET 0x1A00000
@@ -580,4 +583,29 @@ struct edgetpu_ext_mailbox_ioctl {
 #define EDGETPU_GET_FATAL_ERRORS \
 	_IOR(EDGETPU_IOCTL_BASE, 32, __u32)
 
+/*
+ * struct edgetpu_test_ext_ioctl
+ * @fd:			file descriptor of opened TPU device
+ * @mbox_bmap:		bitmap of requested mailboxes
+ * @client_type:	value based on type of client
+ * @cmd:		command to TPU driver on external interface
+ * @attrs:		mailbox attributes (pointer to
+ *			edgetpu_mailbox_attr, can be NULL)
+ */
+struct edgetpu_test_ext_ioctl {
+	__s32 fd;
+	__u32 mbox_bmap;
+	__u32 client_type;
+	__u32 cmd;
+	__u64 attrs;
+};
+
+/*
+ * Invokes external interface used for interoperating with TPU device.
+ *
+ * Usage is only for testing purpose and limited to ROOT user only.
+ */
+#define EDGETPU_TEST_EXTERNAL \
+	_IOW(EDGETPU_IOCTL_BASE, 33, struct edgetpu_test_ext_ioctl)
+
 #endif /* __EDGETPU_H__ */
diff --git a/drivers/edgetpu/mm-backport.h b/drivers/edgetpu/mm-backport.h
index 8831285..ba43f54 100644
--- a/drivers/edgetpu/mm-backport.h
+++ b/drivers/edgetpu/mm-backport.h
@@ -21,6 +21,7 @@
 #define FOLL_LONGTERM 0
 
 #define pin_user_pages_fast get_user_pages_fast
+#define pin_user_pages get_user_pages
 #define unpin_user_page put_page
 
 #ifndef untagged_addr
diff --git a/drivers/edgetpu/mobile-debug-dump.h b/drivers/edgetpu/mobile-debug-dump.h
index 9eaf069..0a9aef9 100644
--- a/drivers/edgetpu/mobile-debug-dump.h
+++ b/drivers/edgetpu/mobile-debug-dump.h
@@ -43,7 +43,7 @@ mobile_sscd_collect_mappings_segment(struct edgetpu_device_group **groups, size_
  *
  * Returns the total number of queues collected since some queues may have been released for groups
  * with detached mailboxes. The return value is less than or equal to the total number of queues
- * expected based on @num_groups i.e (2 * @num_groups +2).
+ * expected based on @num_groups i.e. (2 * @num_groups +2).
  */
 size_t mobile_sscd_collect_cmd_resp_queues(struct edgetpu_dev *etdev,
 					   struct edgetpu_device_group **groups, size_t num_groups,
diff --git a/drivers/edgetpu/mobile-firmware.c b/drivers/edgetpu/mobile-firmware.c
index f021a32..88a2a4d 100644
--- a/drivers/edgetpu/mobile-firmware.c
+++ b/drivers/edgetpu/mobile-firmware.c
@@ -24,23 +24,38 @@
 #include "edgetpu-mobile-platform.h"
 #include "mobile-firmware.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 struct mobile_image_config *mobile_firmware_get_image_config(struct edgetpu_dev *etdev)
 {
-	return (struct mobile_image_config *) edgetpu_firmware_get_data(etdev->firmware);
+	return edgetpu_firmware_get_data(etdev->firmware);
 }
 
-static void mobile_firmware_clear_mappings(struct edgetpu_dev *etdev,
-					   struct mobile_image_config *image_config)
+/* Clear mapping #i */
+static void clear_mapping(struct edgetpu_dev *etdev,
+			  struct mobile_image_config *image_config, int i)
 {
 	tpu_addr_t tpu_addr;
 	size_t size;
+
+	tpu_addr = image_config->mappings[i].virt_address;
+	size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
+	edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
+}
+
+static void mobile_firmware_clear_mappings(struct edgetpu_dev *etdev,
+					   struct mobile_image_config *image_config)
+{
 	int i;
 
-	for (i = 0; i < image_config->num_iommu_mapping; i++) {
-		tpu_addr = image_config->mappings[i].virt_address;
-		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
+	for (i = 0; i < image_config->num_iommu_mapping; i++)
+		clear_mapping(etdev, image_config, i);
 }
 
 static int mobile_firmware_setup_mappings(struct edgetpu_dev *etdev,
@@ -55,6 +70,7 @@ static int mobile_firmware_setup_mappings(struct edgetpu_dev *etdev,
 		tpu_addr = image_config->mappings[i].virt_address;
 		if (!tpu_addr) {
 			etdev_warn(etdev, "Invalid firmware header\n");
+			ret = -EIO;
 			goto err;
 		}
 		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
@@ -76,26 +92,29 @@ static int mobile_firmware_setup_mappings(struct edgetpu_dev *etdev,
 	return 0;
 
 err:
-	while (i--) {
-		tpu_addr = image_config->mappings[i].virt_address;
-		size = CONFIG_TO_SIZE(image_config->mappings[i].image_config_value);
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
+	while (i--)
+		clear_mapping(etdev, image_config, i);
 	return ret;
 }
 
-static void mobile_firmware_clear_ns_mappings(struct edgetpu_dev *etdev,
-					      struct mobile_image_config *image_config)
+static void clear_ns_mapping(struct edgetpu_dev *etdev,
+			     struct mobile_image_config *image_config, int i)
 {
 	tpu_addr_t tpu_addr;
 	size_t size;
+
+	tpu_addr = image_config->ns_iommu_mappings[i] & ~(0xFFF);
+	size = CONFIG_TO_MBSIZE(image_config->ns_iommu_mappings[i]);
+	edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
+}
+
+static void mobile_firmware_clear_ns_mappings(struct edgetpu_dev *etdev,
+					      struct mobile_image_config *image_config)
+{
 	int i;
 
-	for (i = 0; i < image_config->num_ns_iommu_mappings; i++) {
-		tpu_addr = image_config->ns_iommu_mappings[i] & ~(0xFFF);
-		size = CONFIG_TO_MBSIZE(image_config->ns_iommu_mappings[i]);
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
+	for (i = 0; i < image_config->num_ns_iommu_mappings; i++)
+		clear_ns_mapping(etdev, image_config, i);
 }
 
 static int mobile_firmware_setup_ns_mappings(struct edgetpu_dev *etdev,
@@ -129,11 +148,8 @@ static int mobile_firmware_setup_ns_mappings(struct edgetpu_dev *etdev,
 	return 0;
 
 err:
-	while (i--) {
-		tpu_addr = image_config->ns_iommu_mappings[i] & ~(0xFFF);
-		size = CONFIG_TO_MBSIZE(image_config->ns_iommu_mappings[i]);
-		edgetpu_mmu_remove_translation(etdev, tpu_addr, size, EDGETPU_CONTEXT_KCI);
-	}
+	while (i--)
+		clear_ns_mapping(etdev, image_config, i);
 	return ret;
 }
 
@@ -294,6 +310,43 @@ static void program_iremap_csr(struct edgetpu_dev *etdev)
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL + 8, 1);
 }
 
+static void mobile_firmware_setup_ssmt(struct edgetpu_dev *etdev)
+{
+	int i;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct mobile_image_config *image_config = mobile_firmware_get_image_config(etdev);
+
+	/*
+	 * This only works if the SSMT is set to client-driven mode, which only GSA can do.
+	 * Skip if GSA is not available
+	 */
+	if (!etmdev->ssmt_base || !etmdev->gsa_dev)
+		return;
+
+	etdev_dbg(etdev, "Setting up SSMT for privilege level: %d\n",
+		  image_config->privilege_level);
+
+	/*
+	 * Setup non-secure SCIDs, assume VID = SCID when running at TZ or GSA level,
+	 * 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;
+
+		if (image_config->privilege_level == FW_PRIV_LEVEL_NS)
+			val = 0;
+		else
+			val = i;
+
+		writel(val, SSMT_NS_READ_STREAM_VID_REG(etmdev->ssmt_base, i));
+		writel(val, SSMT_NS_WRITE_STREAM_VID_REG(etmdev->ssmt_base, i));
+	}
+}
+
 static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 				       struct edgetpu_firmware_buffer *fw_buf)
 {
@@ -305,6 +358,8 @@ static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 	if (IS_ENABLED(CONFIG_RIO))
 		program_iremap_csr(etdev);
 
+	mobile_firmware_setup_ssmt(etdev);
+
 	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
 }
 
@@ -385,23 +440,24 @@ static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 		goto out;
 	}
 
-	if (image_config_changed) {
-		/* clear last image mappings */
-		if (last_image_config->privilege_level == FW_PRIV_LEVEL_NS)
-			mobile_firmware_clear_mappings(etdev, last_image_config);
+	if (!image_config_changed)
+		goto out;
 
-		if (image_config->privilege_level == FW_PRIV_LEVEL_NS)
-			ret = mobile_firmware_setup_mappings(etdev, image_config);
-		if (ret)
-			goto out;
-		mobile_firmware_clear_ns_mappings(etdev, last_image_config);
-		ret = mobile_firmware_setup_ns_mappings(etdev, image_config);
-		if (ret) {
-			mobile_firmware_clear_mappings(etdev, image_config);
-			goto out;
-		}
-		mobile_firmware_save_image_config(etdev, image_config);
+	/* clear last image mappings */
+	if (last_image_config->privilege_level == FW_PRIV_LEVEL_NS)
+		mobile_firmware_clear_mappings(etdev, last_image_config);
+
+	if (image_config->privilege_level == FW_PRIV_LEVEL_NS)
+		ret = mobile_firmware_setup_mappings(etdev, image_config);
+	if (ret)
+		goto out;
+	mobile_firmware_clear_ns_mappings(etdev, last_image_config);
+	ret = mobile_firmware_setup_ns_mappings(etdev, image_config);
+	if (ret) {
+		mobile_firmware_clear_mappings(etdev, image_config);
+		goto out;
 	}
+	mobile_firmware_save_image_config(etdev, image_config);
 out:
 	memunmap(image_vaddr);
 	return ret;
diff --git a/drivers/edgetpu/mobile-firmware.h b/drivers/edgetpu/mobile-firmware.h
index d17d0c0..27ee048 100644
--- a/drivers/edgetpu/mobile-firmware.h
+++ b/drivers/edgetpu/mobile-firmware.h
@@ -85,5 +85,4 @@ void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev);
  */
 int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset);
 
-
 #endif /* __MOBILE_FIRMWARE_H__ */