Merge remote-tracking branch 'pro' into android-gs-cloudripper-5.10

    * origin/darwinn-2.0: (23 commits)
      edgetpu: remove redundant args
      edgetpu: janeiro: add power management
      edgetpu: janeiro scrub references to TPU CPU model
      edgetpu: abrolhos scrub more TPU CPU references
      edgetpu: remove some references to codenames and hardware details
      edgetpu: remove some details from usage-stats comments and symbols
      edgetpu: fix typo in debug dump header
      Mock components to build TPU driver
      edgetpu: log error when firmware load failed
      edgetpu: google: use default domain when AUX disabled
      edgetpu: fix edgetpu_mmu_alloc_domain memory leak
      edgetpu: remove dependency on iommu_group_id
      edgetpu: janeiro: use NS iommu mapping for f/w
      edgetpu: fix watchdog job cancel ordering
      edgetpu: don't check mailbox_detachable on fs_release
      edgetpu: add dev_inaccessible field to struct group
      edgetpu: hermosa remove irqreturn header from smmu
      edgetpu: remove chip names from mmu.h
      edgetpu: move mailbox disable out from mbox remove
      edgetpu: export mailbox enable and disable functions

Signed-off-by: Nishant Prajapati <nishantpjpt@google.com>
Change-Id: I481b99e32bb07053e9faaa8965259ea163fd4ebc

22 files changed, 609 insertions, 207 deletions

diff --git a/drivers/edgetpu/edgetpu-debug-dump.h b/drivers/edgetpu/edgetpu-debug-dump.h
index ec33668..125ed1a 100644
--- a/drivers/edgetpu/edgetpu-debug-dump.h
+++ b/drivers/edgetpu/edgetpu-debug-dump.h
@@ -44,7 +44,7 @@ struct edgetpu_debug_stats {
 struct edgetpu_dump_segment {
 	u64 type;	/* type of the dump */
 	u64 size;	/* size of the dump data */
-	u64 src_addr; /* source of the dump on the R52 address map */
+	u64 src_addr; /* source of the dump on the CPU address map */
 };
 
 struct edgetpu_debug_dump {
@@ -89,4 +89,4 @@ int edgetpu_get_debug_dump(struct edgetpu_dev *etdev,
  */
 void edgetpu_debug_dump_resp_handler(struct edgetpu_dev *etdev);
 
-#endif /* EDEGETPU_DEBUG_DUMP_H_ */
+#endif /* EDGETPU_DEBUG_DUMP_H_ */
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 3b2e6fc..ae3a8c3 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -156,7 +156,16 @@ static void edgetpu_device_group_kci_leave(struct edgetpu_device_group *group)
 {
 #ifdef EDGETPU_HAS_MULTI_GROUPS
 	edgetpu_kci_update_usage_async(group->etdev);
-	return edgetpu_group_kci_close_device(group);
+	/*
+	 * Theoretically we don't need to check @dev_inaccessible here.
+	 * @dev_inaccessible is true implies the client has wakelock count zero, under such case
+	 * edgetpu_mailbox_deactivate() has been called on releasing the wakelock and therefore this
+	 * edgetpu_group_kci_close_device() call won't send any KCI.
+	 * Still have a check here in case this function does CSR programming other than calling
+	 * edgetpu_mailbox_deactivate() someday.
+	 */
+	if (!group->dev_inaccessible)
+		edgetpu_group_kci_close_device(group);
 #else /* !EDGETPU_HAS_MULTI_GROUPS */
 	struct kci_worker_param *params =
 		kmalloc_array(group->n_clients, sizeof(*params), GFP_KERNEL);
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 3a5e252..545b4bd 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -69,6 +69,22 @@ struct edgetpu_device_group {
 	 * creating this group.
 	 */
 	bool mailbox_detachable;
+	/*
+	 * Whether group->etdev is inaccessible.
+	 * Some group operations will access device CSRs. If the device is known to be
+	 * inaccessible (typically not powered on) then set this field to true to
+	 * prevent HW interactions.
+	 *
+	 * This field is always false for !EDGETPU_HAS_WAKELOCK chipsets.
+	 *
+	 * For EDGETPU_HAS_MCP chipsets this field should be replaced with a
+	 * boolean array with size @n_clients, but we don't have a chipset with
+	 * EDGETPU_HAS_MCP && EDGETPU_HAS_WAKELOCK yet.
+	 *
+	 * Is not protected by @lock because this is only written when releasing the
+	 * leader of this group.
+	 */
+	bool dev_inaccessible;
 
 	/* protects everything in the following comment block */
 	struct mutex lock;
diff --git a/drivers/edgetpu/edgetpu-firmware.c b/drivers/edgetpu/edgetpu-firmware.c
index d0dc575..00da3c4 100644
--- a/drivers/edgetpu/edgetpu-firmware.c
+++ b/drivers/edgetpu/edgetpu-firmware.c
@@ -180,7 +180,7 @@ static int edgetpu_firmware_load_locked(
 	if (handlers && handlers->alloc_buffer) {
 		ret = handlers->alloc_buffer(et_fw, &fw_desc->buf);
 		if (ret) {
-			etdev_dbg(etdev, "handler alloc_buffer failed: %d\n",
+			etdev_err(etdev, "handler alloc_buffer failed: %d\n",
 				  ret);
 			return ret;
 		}
@@ -188,14 +188,14 @@ static int edgetpu_firmware_load_locked(
 
 	ret = edgetpu_firmware_do_load_locked(et_fw, fw_desc, name);
 	if (ret) {
-		etdev_dbg(etdev, "firmware request failed: %d\n", ret);
+		etdev_err(etdev, "firmware request failed: %d\n", ret);
 		goto out_free_buffer;
 	}
 
 	if (handlers && handlers->setup_buffer) {
 		ret = handlers->setup_buffer(et_fw, &fw_desc->buf);
 		if (ret) {
-			etdev_dbg(etdev, "handler setup_buffer failed: %d\n",
+			etdev_err(etdev, "handler setup_buffer failed: %d\n",
 				  ret);
 			goto out_do_unload_locked;
 		}
@@ -467,7 +467,7 @@ int edgetpu_firmware_run_locked(struct edgetpu_firmware *et_fw,
 	}
 
 	/*
-	 * Previous firmware buffer is not used anymore when R52 runs on
+	 * Previous firmware buffer is not used anymore when the CPU runs on
 	 * new firmware buffer. Unload this before et_fw->p->fw_buf is
 	 * overwritten by new buffer information.
 	 */
diff --git a/drivers/edgetpu/edgetpu-firmware.h b/drivers/edgetpu/edgetpu-firmware.h
index e41543d..3b784c5 100644
--- a/drivers/edgetpu/edgetpu-firmware.h
+++ b/drivers/edgetpu/edgetpu-firmware.h
@@ -39,7 +39,7 @@ enum edgetpu_fw_flavor {
 	FW_FLAVOR_BL1 = 1,
 	/* systest app image */
 	FW_FLAVOR_SYSTEST = 2,
-	/* default production app image from DarwiNN team */
+	/* default production app image */
 	FW_FLAVOR_PROD_DEFAULT = 3,
 	/* custom image produced by other teams */
 	FW_FLAVOR_CUSTOM = 4,
@@ -140,7 +140,7 @@ struct edgetpu_firmware_handlers {
 				struct edgetpu_firmware_buffer *fw_buf);
 	/*
 	 * Platform-specific handling after firmware loaded, before running
-	 * the firmware, such as validating the firmware or resetting the R52
+	 * the firmware, such as validating the firmware or resetting the
 	 * processor.
 	 */
 	int (*prepare_run)(struct edgetpu_firmware *et_fw,
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 5b07632..256f946 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -107,16 +107,18 @@ static int edgetpu_fs_release(struct inode *inode, struct file *file)
 	wakelock_count = edgetpu_wakelock_lock(client->wakelock);
 	mutex_lock(&client->group_lock);
 	/*
-	 * @wakelock = 0 means the device might be powered off. And for group with a non-detachable
-	 * mailbox, its mailbox is removed when the group is released, in such case we need to
-	 * ensure the device is powered to prevent kernel panic on programming VII mailbox CSRs.
-	 *
-	 * For mailbox-detachable groups the mailbox had been removed when the wakelock was
-	 * released, edgetpu_device_group_release() doesn't need the device be powered in this case.
+	 * @wakelock_count = 0 means the device might be powered off. Mailbox is removed when the
+	 * group is released, we need to ensure the device is powered to prevent kernel panic on
+	 * programming VII mailbox CSRs.
+	 * If the device is known to be not powered then simply set dev_inaccessible to true to
+	 * prevent device interactions during group releasing.
 	 */
-	if (!wakelock_count && client->group && !client->group->mailbox_detachable) {
-		wakelock_count = 1;
-		edgetpu_pm_get(etdev->pm);
+	if (!wakelock_count && client->group) {
+		/* assumes @group->etdev == @client->etdev, i.e. @client is the leader of @group */
+		if (edgetpu_pm_get_if_powered(etdev->pm))
+			wakelock_count = 1;
+		else
+			client->group->dev_inaccessible = true;
 	}
 	mutex_unlock(&client->group_lock);
 	edgetpu_wakelock_unlock(client->wakelock);
@@ -933,8 +935,10 @@ static void edgetpu_fs_setup_debugfs(struct edgetpu_dev *etdev)
 {
 	etdev->d_entry =
 		debugfs_create_dir(etdev->dev_name, edgetpu_debugfs_dir);
-	if (!etdev->d_entry)
+	if (IS_ERR_OR_NULL(etdev->d_entry)) {
+		etdev_warn(etdev, "Failed to setup debugfs\n");
 		return;
+	}
 	debugfs_create_file("mappings", 0440, etdev->d_entry,
 			    etdev, &mappings_ops);
 	debugfs_create_file("statusregs", 0440, etdev->d_entry, etdev,
@@ -1036,18 +1040,16 @@ static const struct file_operations syncfences_ops = {
 	.release = single_release,
 };
 
-static int edgetpu_debugfs_global_setup(void)
+static void edgetpu_debugfs_global_setup(void)
 {
 	edgetpu_debugfs_dir = debugfs_create_dir("edgetpu", NULL);
-	if (IS_ERR(edgetpu_debugfs_dir)) {
-		pr_err(DRIVER_NAME " error creating edgetpu debugfs dir: %ld\n",
-		       PTR_ERR(edgetpu_debugfs_dir));
-		return PTR_ERR(edgetpu_debugfs_dir);
+	if (IS_ERR_OR_NULL(edgetpu_debugfs_dir)) {
+		pr_warn(DRIVER_NAME " error creating edgetpu debugfs dir\n");
+		return;
 	}
 
 	debugfs_create_file("syncfences", 0440, edgetpu_debugfs_dir, NULL,
 			    &syncfences_ops);
-	return 0;
 }
 
 int __init edgetpu_fs_init(void)
diff --git a/drivers/edgetpu/edgetpu-google-iommu.c b/drivers/edgetpu/edgetpu-google-iommu.c
index cdd01d9..9d28949 100644
--- a/drivers/edgetpu/edgetpu-google-iommu.c
+++ b/drivers/edgetpu/edgetpu-google-iommu.c
@@ -71,6 +71,9 @@ get_domain_by_context_id(struct edgetpu_dev *etdev,
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	uint pasid;
 
+	/* always return the default domain when AUX is not supported */
+	if (!etiommu->aux_enabled)
+		return iommu_get_domain_for_dev(dev);
 	if (ctx_id == EDGETPU_CONTEXT_INVALID)
 		return NULL;
 	if (ctx_id & EDGETPU_CONTEXT_DOMAIN_TOKEN)
@@ -240,13 +243,10 @@ int edgetpu_mmu_attach(struct edgetpu_dev *etdev, void *mmu_info)
 	idr_init(&etiommu->domain_pool);
 	mutex_init(&etiommu->pool_lock);
 	etiommu->iommu_group = iommu_group_get(etdev->dev);
-	if (etiommu->iommu_group) {
+	if (etiommu->iommu_group)
 		iommu_group_set_name(etiommu->iommu_group, "edgetpu");
-		dev_dbg(etdev->dev, "iommu group id %d setup\n",
-			iommu_group_id(etiommu->iommu_group));
-	} else {
+	else
 		dev_warn(etdev->dev, "device has no iommu group\n");
-	}
 
 	iommu_dev_enable_feature(etdev->dev, IOMMU_DEV_FEAT_AUX);
 	if (!iommu_dev_feature_enabled(etdev->dev, IOMMU_DEV_FEAT_AUX))
@@ -630,14 +630,11 @@ static struct edgetpu_iommu_domain invalid_etdomain = {
 
 struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_iommu_domain *etdomain =
-		kzalloc(sizeof(*etdomain), GFP_KERNEL);
+	struct edgetpu_iommu_domain *etdomain;
 	struct edgetpu_iommu *etiommu = etdev->mmu_cookie;
 	struct iommu_domain *domain;
 	int token;
 
-	if (!etdomain)
-		return NULL;
 	if (!etiommu->aux_enabled)
 		return &invalid_etdomain;
 	domain = iommu_domain_alloc(etdev->dev->bus);
@@ -646,15 +643,23 @@ struct edgetpu_iommu_domain *edgetpu_mmu_alloc_domain(struct edgetpu_dev *etdev)
 		return NULL;
 	}
 
+	etdomain = kzalloc(sizeof(*etdomain), GFP_KERNEL);
+	if (!etdomain) {
+		iommu_domain_free(domain);
+		return NULL;
+	}
+
 	mutex_lock(&etiommu->pool_lock);
 	token = idr_alloc(&etiommu->domain_pool, domain, 0,
 			  EDGETPU_DOMAIN_TOKEN_END, GFP_KERNEL);
 	mutex_unlock(&etiommu->pool_lock);
 	if (token < 0) {
 		etdev_warn(etdev, "alloc iommu domain token failed: %d", token);
+		kfree(etdomain);
 		iommu_domain_free(domain);
 		return NULL;
 	}
+
 	edgetpu_init_etdomain(etdomain, domain, token);
 	return etdomain;
 }
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index 1258cf0..f2f2ff1 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -416,7 +416,7 @@ int edgetpu_get_state_errno_locked(struct edgetpu_dev *etdev);
 
 /*
  * "External mailboxes" below refers to mailboxes that are not handled
- * directly by the DarwiNN runtime, such as secure or device-to-device.
+ * directly by the runtime, such as secure or device-to-device.
  *
  * Chip specific code will typically keep track of state and inform the firmware
  * that a mailbox has become active/inactive.
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index e467fac..c62ac73 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -487,7 +487,7 @@ int edgetpu_kci_init(struct edgetpu_mailbox_manager *mgr,
 	INIT_WORK(&kci->work, edgetpu_kci_consume_responses_work);
 	edgetpu_reverse_kci_init(&kci->rkci);
 	INIT_WORK(&kci->usage_work, edgetpu_kci_update_usage_work);
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+	edgetpu_mailbox_enable(mailbox);
 	return 0;
 }
 
@@ -509,7 +509,7 @@ int edgetpu_kci_reinit(struct edgetpu_kci *kci)
 	if (ret)
 		return ret;
 	edgetpu_mailbox_init_doorbells(mailbox);
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+	edgetpu_mailbox_enable(mailbox);
 
 	return 0;
 }
diff --git a/drivers/edgetpu/edgetpu-mailbox.c b/drivers/edgetpu/edgetpu-mailbox.c
index 11b6fe2..1a19185 100644
--- a/drivers/edgetpu/edgetpu-mailbox.c
+++ b/drivers/edgetpu/edgetpu-mailbox.c
@@ -80,12 +80,6 @@ edgetpu_mailbox_create_locked(struct edgetpu_mailbox_manager *mgr, uint index)
 	return mailbox;
 }
 
-/* Disables a mailbox via setting CSR. */
-static void edgetpu_mailbox_disable(struct edgetpu_mailbox *mailbox)
-{
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 0);
-}
-
 /*
  * Disables the @index-th mailbox via setting CSR. Doesn't need
  * @mgr->mailboxes[index] be allocated.
@@ -203,12 +197,12 @@ int edgetpu_mailbox_set_queue(struct edgetpu_mailbox *mailbox,
 /* Reset mailbox queues, clear out any commands/responses left from before. */
 void edgetpu_mailbox_reset(struct edgetpu_mailbox *mailbox)
 {
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 0);
+	edgetpu_mailbox_disable(mailbox);
 	EDGETPU_MAILBOX_CMD_QUEUE_WRITE(mailbox, head, 0);
 	edgetpu_mailbox_set_cmd_queue_tail(mailbox, 0);
 	edgetpu_mailbox_set_resp_queue_head(mailbox, 0);
 	EDGETPU_MAILBOX_RESP_QUEUE_WRITE(mailbox, tail, 0);
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+	edgetpu_mailbox_enable(mailbox);
 }
 
 /* Sets the priority of @mailbox. */
@@ -263,6 +257,8 @@ edgetpu_mailbox_vii_add(struct edgetpu_mailbox_manager *mgr, uint id)
  *
  * Returns 0 on success, or a negative errno on error.
  * Returns -EBUSY if any mailbox is using.
+ *
+ * Caller calls edgetpu_mailbox_enable() to enable the returned mailboxes.
  */
 int edgetpu_mailbox_p2p_batch(struct edgetpu_mailbox_manager *mgr, uint n,
 			      uint skip_i, struct edgetpu_mailbox **mailboxes)
@@ -343,8 +339,7 @@ out:
  * Removes a mailbox from the manager.
  * Returns 0 on success.
  */
-int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr,
-			   struct edgetpu_mailbox *mailbox)
+int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr, struct edgetpu_mailbox *mailbox)
 {
 	unsigned long flags;
 
@@ -357,7 +352,6 @@ int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr,
 	}
 
 	mgr->mailboxes[mailbox->mailbox_id] = NULL;
-	edgetpu_mailbox_disable(mailbox);
 	/* KCI mailbox is a special case */
 	if (mailbox->mailbox_id == KERNEL_MAILBOX_INDEX)
 		edgetpu_kci_release(mgr->etdev, mailbox->internal.kci);
@@ -467,7 +461,7 @@ int edgetpu_mailbox_init_vii(struct edgetpu_vii *vii,
 	mailbox->internal.group = edgetpu_device_group_get(group);
 	vii->etdev = group->etdev;
 	vii->mailbox = mailbox;
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+	edgetpu_mailbox_enable(mailbox);
 	return 0;
 }
 
@@ -479,6 +473,8 @@ void edgetpu_mailbox_remove_vii(struct edgetpu_vii *vii)
 	edgetpu_mailbox_free_queue(etdev, vii->mailbox, &vii->cmd_queue_mem);
 	edgetpu_mailbox_free_queue(etdev, vii->mailbox, &vii->resp_queue_mem);
 	if (vii->mailbox) {
+		if (!vii->mailbox->internal.group->dev_inaccessible)
+			edgetpu_mailbox_disable(vii->mailbox);
 		edgetpu_device_group_put(vii->mailbox->internal.group);
 		edgetpu_mailbox_remove(etdev->mailbox_manager, vii->mailbox);
 		vii->mailbox = NULL;
@@ -704,7 +700,7 @@ void edgetpu_mailbox_reinit_vii(struct edgetpu_device_group *group)
 	edgetpu_mailbox_set_queue(mailbox, MAILBOX_RESP_QUEUE,
 				  group->vii.resp_queue_mem.tpu_addr,
 				  resp_queue_size);
-	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+	edgetpu_mailbox_enable(mailbox);
 }
 
 void edgetpu_mailbox_restore_active_vii_queues(struct edgetpu_dev *etdev)
diff --git a/drivers/edgetpu/edgetpu-mailbox.h b/drivers/edgetpu/edgetpu-mailbox.h
index c212a8a..8425807 100644
--- a/drivers/edgetpu/edgetpu-mailbox.h
+++ b/drivers/edgetpu/edgetpu-mailbox.h
@@ -192,10 +192,13 @@ edgetpu_mailbox_create_mgr(struct edgetpu_dev *etdev,
 /* interrupt handler */
 irqreturn_t edgetpu_mailbox_handle_irq(struct edgetpu_mailbox_manager *mgr);
 
-/* removes the mailbox previously requested from a mailbox manager */
-int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr,
-			   struct edgetpu_mailbox *mailbox);
-/* removes all the mailboxes previously requested */
+/*
+ * Removes the mailbox previously requested from a mailbox manager.
+ *
+ * This function doesn't change the state of mailbox enable/disable.
+ */
+int edgetpu_mailbox_remove(struct edgetpu_mailbox_manager *mgr, struct edgetpu_mailbox *mailbox);
+/* Removes and disables all the mailboxes previously requested. */
 void edgetpu_mailbox_remove_all(struct edgetpu_mailbox_manager *mgr);
 
 /* configure mailbox */
@@ -231,11 +234,14 @@ void edgetpu_mailbox_inc_resp_queue_head(struct edgetpu_mailbox *mailbox,
  * Request the mailbox with mailbox_id equals @id.
  * @id = 0 means there is no preference, @mgr will return a spare mailbox.
  *
+ * Caller calls edgetpu_mailbox_enable() to enable the returned mailbox.
+ *
  * -EBUSY is returned if the requested @id is used or there is no mailbox
  * available.
  */
 struct edgetpu_mailbox *
 edgetpu_mailbox_vii_add(struct edgetpu_mailbox_manager *mgr, uint id);
+
 /*
  * Validates the mailbox attributes.
  * Returns 0 if valid, otherwise a negative errno.
@@ -425,4 +431,16 @@ static inline u32 circular_queue_inc(u32 index, u32 inc, u32 queue_size)
 			      struct edgetpu_mailbox_resp_queue_csr, field, \
 			      value)
 
+/* Enables a mailbox by setting CSR. */
+static inline void edgetpu_mailbox_enable(struct edgetpu_mailbox *mailbox)
+{
+	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 1);
+}
+
+/* Disables a mailbox by setting CSR. */
+static inline void edgetpu_mailbox_disable(struct edgetpu_mailbox *mailbox)
+{
+	EDGETPU_MAILBOX_CONTEXT_WRITE(mailbox, context_enable, 0);
+}
+
 #endif /* __EDGETPU_MAILBOX_H__ */
diff --git a/drivers/edgetpu/edgetpu-mmu.h b/drivers/edgetpu/edgetpu-mmu.h
index 8c5ae3c..094f14d 100644
--- a/drivers/edgetpu/edgetpu-mmu.h
+++ b/drivers/edgetpu/edgetpu-mmu.h
@@ -199,9 +199,7 @@ void edgetpu_mmu_free(struct edgetpu_dev *etdev, tpu_addr_t tpu_addr,
  *
  * Description: Add a mapping from iova -> paddr to the MMU for the chip.
  * paddr can be considered a physical address from the TPU's viewpoint, but
- * may actually be another IOVA for another IOMMU downstream of the chip MMU
- * (as on Hermosa, where the SMMU translates TPU VAs to IOVAs sent to the IOMMU
- * downstream of the TPU).
+ * may actually be another IOVA for another IOMMU downstream of the chip MMU.
  *
  * Note: for chipsets with edgetpu_mmu_alloc() support, @iova passed to this
  * function must be either allocated from edgetpu_mmu_alloc() or reserved by
@@ -230,12 +228,12 @@ void edgetpu_mmu_remove_translation(struct edgetpu_dev *etdev,
  * @context_id: context ID for the mapping
  * @mmu_flags: the flag or'ed with EDGETPU_MMU_* macros
  *
- * Description: For chips with internal MMUs (e.g., Hermosa SMMU), add the
- * required internal MMU mapping for the TPU to access @downstream_addr, the
- * DMA or physical address of the buffer as returned by the Linux DMA API when
- * the DMA mapping was created.  This can be used with, for example, buffers
- * allocated using dma_alloc_coherent(), which are mapped appropriately for
- * any downstream IOMMU and must be mapped to the TPU internal MMU as well.
+ * Description: For chips with internal MMUs, add the required internal MMU
+ * mapping for the TPU to access @down_addr, the DMA or physical address of the
+ * buffer as returned by the Linux DMA API when the DMA mapping was created.
+ * This can be used with, for example, buffers allocated using
+ * dma_alloc_coherent(), which are mapped appropriately for any downstream IOMMU
+ * and must be mapped to the TPU internal MMU as well.
  *
  * For a chip that doesn't have an internal MMU but has the IOMMU domain AUX
  * feature, perform the necessary mapping to @context_id and return the
@@ -261,9 +259,8 @@ void edgetpu_mmu_tpu_unmap(struct edgetpu_dev *etdev,
  * @context_id: context ID for the mapping
  * @mmu_flags: the flag or'ed with EDGETPU_MMU_* macros
  *
- * Description: For chips with internal MMUs (e.g., Hermosa SMMU), add the
- * required internal MMU mapping for the TPU to access the DMA addresses of
- * @sgt.
+ * Description: For chips with internal MMUs, add the required internal MMU
+ * mapping for the TPU to access the DMA addresses of @sgt.
  *
  * For a chip that doesn't have an internal MMU but has the IOMMU domain AUX
  * feature, perform the necessary mapping to @context_id and return the
diff --git a/drivers/edgetpu/edgetpu-pm.c b/drivers/edgetpu/edgetpu-pm.c
index 1e28141..df1c179 100644
--- a/drivers/edgetpu/edgetpu-pm.c
+++ b/drivers/edgetpu/edgetpu-pm.c
@@ -242,7 +242,7 @@ static int pchannel_state_change_request(struct edgetpu_dev *etdev, int state)
 	if (state == STATE_RUN)
 		return 0;
 
-	/* Phase 3: R52 acknowledgment */
+	/* Phase 3: CPU acknowledgment */
 	ret = etdev_poll_power_state(etdev, val,
 				     (val & PACCEPT) || (val & PDENY));
 	if (val & PDENY) {
diff --git a/drivers/edgetpu/edgetpu-sw-watchdog.c b/drivers/edgetpu/edgetpu-sw-watchdog.c
index 5d96e4d..4e7f681 100644
--- a/drivers/edgetpu/edgetpu-sw-watchdog.c
+++ b/drivers/edgetpu/edgetpu-sw-watchdog.c
@@ -5,6 +5,7 @@
  * Copyright (C) 2020 Google, Inc.
  */
 
+#include <asm/barrier.h>
 #include <linux/atomic.h>
 #include <linux/module.h>
 #include <linux/slab.h>
@@ -124,21 +125,29 @@ int edgetpu_sw_wdt_create(struct edgetpu_dev *etdev, unsigned long active_ms,
 
 int edgetpu_sw_wdt_start(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_sw_wdt *etdev_sw_wdt = etdev->etdev_sw_wdt;
+	struct edgetpu_sw_wdt *wdt;
 
-	if (!etdev_sw_wdt)
+	/* to match edgetpu_sw_wdt_destroy() */
+	smp_mb();
+	wdt = etdev->etdev_sw_wdt;
+	if (!wdt)
 		return -EINVAL;
-	if (!etdev_sw_wdt->et_action_work.edgetpu_sw_wdt_handler)
+	if (!wdt->et_action_work.edgetpu_sw_wdt_handler)
 		etdev_err(etdev, "sw wdt handler not set\n");
-	sw_wdt_start(etdev_sw_wdt);
+	sw_wdt_start(wdt);
 	return 0;
 }
 
 void edgetpu_sw_wdt_stop(struct edgetpu_dev *etdev)
 {
-	if (!etdev->etdev_sw_wdt)
+	struct edgetpu_sw_wdt *wdt;
+
+	/* to match edgetpu_sw_wdt_destroy() */
+	smp_mb();
+	wdt = etdev->etdev_sw_wdt;
+	if (!wdt)
 		return;
-	sw_wdt_stop(etdev->etdev_sw_wdt);
+	sw_wdt_stop(wdt);
 }
 
 void edgetpu_sw_wdt_destroy(struct edgetpu_dev *etdev)
@@ -149,9 +158,14 @@ void edgetpu_sw_wdt_destroy(struct edgetpu_dev *etdev)
 	if (!wdt)
 		return;
 	etdev->etdev_sw_wdt = NULL;
+	/*
+	 * To ensure that etdev->etdev_sw_wdt is NULL so wdt_start() calls from other processes
+	 * won't start the watchdog again.
+	 */
+	smp_mb();
+	sw_wdt_stop(wdt);
 	/* cancel and sync work due to watchdog bite to prevent UAF */
 	cancel_work_sync(&wdt->et_action_work.work);
-	sw_wdt_stop(wdt);
 	counter = atomic_read(&wdt->active_counter);
 	if (counter)
 		etdev_warn(etdev, "Unbalanced WDT active counter: %d", counter);
diff --git a/drivers/edgetpu/edgetpu-usage-stats.h b/drivers/edgetpu/edgetpu-usage-stats.h
index 7ea3e9d..20d5ad7 100644
--- a/drivers/edgetpu/edgetpu-usage-stats.h
+++ b/drivers/edgetpu/edgetpu-usage-stats.h
@@ -36,9 +36,9 @@ struct tpu_usage {
  * Must be kept in sync with firmware struct Component.
  */
 enum edgetpu_usage_component {
-	/* The device as a whole (TPU, R52, DMA330, etc.) */
+	/* The device as a whole */
 	EDGETPU_USAGE_COMPONENT_DEVICE = 0,
-	/* Just the TPU core (scalar core and tiles) */
+	/* Just the TPU core */
 	EDGETPU_USAGE_COMPONENT_TPU = 1,
 	EDGETPU_USAGE_COMPONENT_COUNT = 2, /* number of components above */
 };
@@ -112,20 +112,9 @@ struct __packed edgetpu_usage_max_watermark {
 /* An enum to identify the tracked firmware threads. */
 /* Must be kept in sync with firmware enum class UsageTrackerThreadId. */
 enum edgetpu_usage_threadid {
-	EDGETPU_FW_THREAD_MAIN = 0,
-	EDGETPU_FW_THREAD_KCI_HANDLER = 1,
-	EDGETPU_FW_THREAD_POWER_ADMIN = 2,
-	EDGETPU_FW_THREAD_VII_SCHEDULER = 3,
-	EDGETPU_FW_THREAD_VII_HANDLER = 4,
-	EDGETPU_FW_THREAD_MCP_GRAPH_DRIVER = 5,
-	EDGETPU_FW_THREAD_SCP_GRAPH_DRIVER = 6,
-	EDGETPU_FW_THREAD_TPU_DRIVER = 7,
-	EDGETPU_FW_THREAD_RESTART_HANDLER = 8,
-	EDGETPU_FW_THREAD_POLL_SERVICE = 9,
-	EDGETPU_FW_THREAD_DMA_DRIVER = 10,
-	EDGETPU_FW_THREAD_GRAPH_DMA_DRIVER = 11,
-
-	/* Number of task identifiers above. */
+	/* Individual thread IDs are not tracked. */
+
+	/* Number of task identifiers. */
 	EDGETPU_FW_THREAD_COUNT = 12,
 };
 
diff --git a/drivers/edgetpu/janeiro-device.c b/drivers/edgetpu/janeiro-device.c
index 7016de3..00d480a 100644
--- a/drivers/edgetpu/janeiro-device.c
+++ b/drivers/edgetpu/janeiro-device.c
@@ -12,6 +12,7 @@
 #include "edgetpu-mailbox.h"
 #include "edgetpu-telemetry.h"
 #include "janeiro-platform.h"
+#include "janeiro-pm.h"
 
 static irqreturn_t janeiro_mailbox_handle_irq(struct edgetpu_dev *etdev,
 					      int irq)
@@ -83,11 +84,29 @@ struct edgetpu_dumpregs_range edgetpu_chip_tile_statusregs_ranges[] = {
 int edgetpu_chip_tile_statusregs_nranges =
 	ARRAY_SIZE(edgetpu_chip_tile_statusregs_ranges);
 
+static void edgetpu_chip_set_pm_qos(struct edgetpu_dev *etdev, u32 value)
+{
+}
+
+static void edgetpu_chip_set_bts(struct edgetpu_dev *etdev, u32 value)
+{
+}
+
 void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
 				    struct edgetpu_kci_response_element *resp)
 {
-	etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n", __func__,
-		   resp->code);
+	switch (resp->code) {
+	case RKCI_CODE_PM_QOS:
+		edgetpu_chip_set_pm_qos(etdev, resp->retval);
+		break;
+	case RKCI_CODE_BTS:
+		edgetpu_chip_set_bts(etdev, resp->retval);
+		break;
+	default:
+		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n",
+			   __func__, resp->code);
+		break;
+	}
 }
 
 
diff --git a/drivers/edgetpu/janeiro-firmware.c b/drivers/edgetpu/janeiro-firmware.c
index 966d539..4f0ce84 100644
--- a/drivers/edgetpu/janeiro-firmware.c
+++ b/drivers/edgetpu/janeiro-firmware.c
@@ -53,10 +53,10 @@ struct janeiro_firmware_data {
 	struct iommu_mapping mappings[MAX_IOMMU_MAPPINGS];
 };
 /*
- * Sets the reset state of the R52 core.
+ * Sets the reset state of the TPU CPU.
  * @val: 1 to put the core in reset state, 0 to release core from reset state.
  */
-static void r52_reset(struct edgetpu_dev *etdev, u64 val)
+static void tpu_cpu_reset(struct edgetpu_dev *etdev, u64 val)
 {
 	edgetpu_dev_write_32_sync(etdev, EDGETPU_REG_RESET_CONTROL, val);
 }
@@ -79,6 +79,7 @@ static void janeiro_firmware_before_destroy(struct edgetpu_firmware *et_fw)
 	u32 i, tpu_addr, size;
 	struct edgetpu_dev *etdev = et_fw->etdev;
 
+	tpu_cpu_reset(et_fw->etdev, 1);
 	/* TODO(b/189906347): Remove when GSA/TZ support is available. */
 	/* Remove mappings created by setup_buffer() */
 	data = edgetpu_firmware_get_data(et_fw);
@@ -93,7 +94,6 @@ static void janeiro_firmware_before_destroy(struct edgetpu_firmware *et_fw)
 		edgetpu_firmware_set_data(et_fw, NULL);
 		kfree(data);
 	}
-	r52_reset(et_fw->etdev, 1);
 }
 
 static int janeiro_firmware_alloc_buffer(struct edgetpu_firmware *et_fw,
@@ -237,16 +237,12 @@ static int janeiro_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 //				   fw_buf->used_size - JANEIRO_FW_HEADER_SIZE,
 //				   DMA_TO_DEVICE);
 #endif
-	r52_reset(etdev, 1);
+	tpu_cpu_reset(etdev, 1);
 
 	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
 	edgetpu_mailbox_reset(etdev->kci->mailbox);
 
-	/* Remap TPU CPU instructions to the carveout IOVA. */
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE,
-			     /*FW_IOVA*/ fw_dma_addr);
-	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL, 1);
-	r52_reset(etdev, 0);
+	tpu_cpu_reset(etdev, 0);
 	//TODO: cleanup
 	return ret;
 }
diff --git a/drivers/edgetpu/janeiro-platform.c b/drivers/edgetpu/janeiro-platform.c
index 96f8b60..d330f48 100644
--- a/drivers/edgetpu/janeiro-platform.c
+++ b/drivers/edgetpu/janeiro-platform.c
@@ -13,6 +13,7 @@
 #include <linux/platform_device.h>
 
 #include "edgetpu-config.h"
+#include "edgetpu-debug-dump.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
@@ -100,52 +101,6 @@ janeiro_platform_cleanup_fw_region(struct janeiro_platform_dev *etpdev)
 	etpdev->shared_mem_vaddr = NULL;
 }
 
-static int
-janeiro_platform_map_reserved_region(struct janeiro_platform_dev *etpdev)
-{
-	int ret;
-	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
-
-	ret = edgetpu_mmu_add_translation(etdev, etpdev->fw_region_paddr,
-		etpdev->fw_region_paddr, etpdev->fw_region_size,
-		IOMMU_READ | IOMMU_WRITE, EDGETPU_CONTEXT_KCI);
-	if (ret) {
-		dev_err(etdev->dev,
-			"Unable to map reserved area for firmware\n");
-		return ret;
-	}
-	ret = edgetpu_mmu_add_translation(etdev,
-			etpdev->fw_region_paddr + etpdev->fw_region_size,
-			etpdev->fw_region_paddr + etpdev->fw_region_size,
-			EDGETPU_REMAPPED_DATA_SIZE, IOMMU_READ | IOMMU_WRITE,
-			EDGETPU_CONTEXT_KCI);
-	if (ret) {
-		dev_err(etdev->dev,
-			"Unable to map reserved area for data\n");
-		edgetpu_mmu_remove_translation(etdev,
-					       etpdev->fw_region_paddr,
-					       etpdev->fw_region_size,
-					       EDGETPU_CONTEXT_KCI);
-		return ret;
-	}
-	return 0;
-}
-
-static void
-janeiro_platform_unmap_reserved_region(struct janeiro_platform_dev *etpdev)
-{
-	struct edgetpu_dev *etdev = &etpdev->edgetpu_dev;
-
-	edgetpu_mmu_remove_translation(etdev,
-			etpdev->fw_region_paddr + etpdev->fw_region_size,
-			EDGETPU_REMAPPED_DATA_SIZE,
-			EDGETPU_CONTEXT_KCI);
-	edgetpu_mmu_remove_translation(etdev,
-			etpdev->fw_region_paddr,
-			etpdev->fw_region_size,
-			EDGETPU_CONTEXT_KCI);
-}
-
 int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev)
 {
 	int ret;
@@ -161,41 +116,6 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
-#define EDGETPU_PSM0_CFG 0x1c1880
-#define EDGETPU_PSM0_START 0x1c1884
-#define EDGETPU_PSM0_STATUS 0x1c1888
-#define EDGETPU_PSM1_CFG 0x1c2880
-#define EDGETPU_PSM1_START 0x1c2884
-#define EDGETPU_PSM1_STATUS 0x1c2888
-//TODO: set timeout lower for silicon
-#define EDGETPU_LPM_CHANGE_TIMEOUT	30000
-
-static int janeiro_set_lpm(struct edgetpu_dev *etdev)
-{
-	int ret;
-	u32 val;
-
-	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_START, 1);
-	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM0_STATUS, val,
-				 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, 5, EDGETPU_LPM_CHANGE_TIMEOUT);
-	if (ret) {
-		etdev_err(etdev, "Set LPM1 failed: %d\n", ret);
-		return ret;
-	}
-
-	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_CFG, 0);
-	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM1_CFG, 0);
-
-	return 0;
-}
-
 /*
  * Set shareability for enabling IO coherency in Janeiro
  */
@@ -311,14 +231,6 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 		goto out_remove_device;
 	}
 
-	ret = janeiro_set_lpm(&edgetpu_pdev->edgetpu_dev);
-	if (ret)
-		dev_err(dev, "LPM init failed: %d\n", ret);
-
-	ret = janeiro_platform_map_reserved_region(edgetpu_pdev);
-	if (ret)
-		goto out_remove_device;
-
 	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_LOG,
 				  &edgetpu_pdev->log_mem);
 	janeiro_get_telemetry_mem(edgetpu_pdev, EDGETPU_TELEMETRY_TRACE,
@@ -328,7 +240,7 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 				     &edgetpu_pdev->log_mem,
 				     &edgetpu_pdev->trace_mem);
 	if (ret)
-		goto out_unmap_reserved;
+		goto out_remove_device;
 
 	ret = mobile_edgetpu_firmware_create(&edgetpu_pdev->edgetpu_dev);
 	if (ret) {
@@ -340,18 +252,13 @@ static int edgetpu_platform_probe(struct platform_device *pdev)
 	dev_info(dev, "%s edgetpu initialized. Build: %s\n",
 		 edgetpu_pdev->edgetpu_dev.dev_name, GIT_REPO_TAG);
 
-	ret = edgetpu_chip_firmware_run(&edgetpu_pdev->edgetpu_dev,
-					EDGETPU_DEFAULT_FIRMWARE_NAME, 0);
-	if (ret)
-		dev_err(dev, "failed to run firmware: %d\n", ret);
+	edgetpu_pm_shutdown(&edgetpu_pdev->edgetpu_dev, false);
 out:
 	dev_dbg(dev, "Probe finished\n");
 
 	return 0;
 out_tel_exit:
 	edgetpu_telemetry_exit(&edgetpu_pdev->edgetpu_dev);
-out_unmap_reserved:
-	janeiro_platform_unmap_reserved_region(edgetpu_pdev);
 out_remove_device:
 	edgetpu_device_remove(&edgetpu_pdev->edgetpu_dev);
 out_destroy_iremap:
@@ -368,14 +275,26 @@ static int edgetpu_platform_remove(struct platform_device *pdev)
 	struct edgetpu_dev *etdev = platform_get_drvdata(pdev);
 	struct janeiro_platform_dev *janeiro_pdev = to_janeiro_dev(etdev);
 
-	mobile_edgetpu_firmware_destroy(etdev);
+	/* TODO(b/189906347): Use edgetpu_device_remove() for cleanup after
+	 * having GSA/TZ support.
+	 */
+	edgetpu_pm_get(etdev->pm);
+
 	for (i = 0; i < EDGETPU_NCONTEXTS; i++) {
 		if (janeiro_pdev->irq[i] >= 0)
 			edgetpu_unregister_irq(etdev, janeiro_pdev->irq[i]);
 	}
 
-	janeiro_platform_unmap_reserved_region(janeiro_pdev);
-	edgetpu_device_remove(etdev);
+	edgetpu_telemetry_exit(etdev);
+	edgetpu_chip_exit(etdev);
+	edgetpu_debug_dump_exit(etdev);
+	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
+	mobile_edgetpu_firmware_destroy(etdev);
+	edgetpu_pm_put(etdev->pm);
+	edgetpu_pm_shutdown(etdev, true);
+	edgetpu_usage_stats_exit(etdev);
+	edgetpu_chip_remove_mmu(etdev);
+	edgetpu_fs_remove(etdev);
 	edgetpu_iremap_pool_destroy(etdev);
 	janeiro_platform_cleanup_fw_region(janeiro_pdev);
 	janeiro_pm_destroy(etdev);
diff --git a/drivers/edgetpu/janeiro-platform.h b/drivers/edgetpu/janeiro-platform.h
index 4e1b7c6..196bf75 100644
--- a/drivers/edgetpu/janeiro-platform.h
+++ b/drivers/edgetpu/janeiro-platform.h
@@ -9,15 +9,25 @@
 
 #include <linux/device.h>
 #include <linux/io.h>
+#include <linux/mutex.h>
 #include <linux/types.h>
 
 #include "edgetpu-internal.h"
+#include "janeiro-pm.h"
 
 #define to_janeiro_dev(etdev)                                                  \
 	container_of(etdev, struct janeiro_platform_dev, edgetpu_dev)
 
+// TODO(b/176881607): merge with abrolhos
+struct janeiro_platform_pwr {
+	struct mutex state_lock;
+	u64 min_state;
+	u64 requested_state;
+};
+
 struct janeiro_platform_dev {
 	struct edgetpu_dev edgetpu_dev;
+	struct janeiro_platform_pwr platform_pwr;
 	int irq[EDGETPU_NCONTEXTS];
 	phys_addr_t fw_region_paddr;
 	void *fw_region_vaddr;
diff --git a/drivers/edgetpu/janeiro-pm.c b/drivers/edgetpu/janeiro-pm.c
index 4302e81..7967d16 100644
--- a/drivers/edgetpu/janeiro-pm.c
+++ b/drivers/edgetpu/janeiro-pm.c
@@ -2,45 +2,392 @@
 /*
  * Janeiro EdgeTPU power management support
  *
- * Copyright (C) 2020 Google, Inc,
+ * Copyright (C) 2020 Google, Inc.
  */
 
+#include <linux/atomic.h>
+#include <linux/delay.h>
+#include <linux/iopoll.h>
+#include <linux/module.h>
 #include <linux/pm_runtime.h>
+#include <linux/version.h>
 
+#include "edgetpu-config.h"
+#include "edgetpu-firmware.h"
 #include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
+#include "edgetpu-mailbox.h"
 #include "edgetpu-pm.h"
+#include "edgetpu-telemetry.h"
+#include "janeiro-platform.h"
 #include "janeiro-pm.h"
 
 #include "edgetpu-pm.c"
 
+/* Default power state */
+static int power_state = TPU_ACTIVE_NOM;
+
+module_param(power_state, int, 0660);
+
+static struct dentry *janeiro_pwr_debugfs_dir;
+
 static int janeiro_pwr_state_init(struct device *dev)
 {
 	int ret;
+	int curr_state;
 
 	pm_runtime_enable(dev);
-	ret = pm_runtime_get_sync(dev);
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+
+	if (curr_state > TPU_OFF) {
+		ret = pm_runtime_get_sync(dev);
+		if (ret) {
+			dev_err(dev, "%s: pm_runtime_get_sync err: %d\n",
+				__func__, ret);
+			return ret;
+		}
+	}
+	ret = exynos_acpm_set_init_freq(TPU_ACPM_DOMAIN, curr_state);
 	if (ret) {
-		dev_err(dev, "pm_runtime_get_sync returned %d\n", ret);
+		dev_err(dev, "error initializing tpu ACPM freq: %d\n", ret);
+		if (curr_state > TPU_OFF)
+			pm_runtime_put_sync(dev);
 		return ret;
 	}
 	return ret;
 }
 
+static int janeiro_pwr_state_set_locked(void *data, u64 val)
+{
+	int ret;
+	int curr_state;
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct device *dev = etdev->dev;
+
+	curr_state = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+
+	dev_dbg(dev, "Power state %d -> %llu\n", curr_state, val);
+
+	if (curr_state == TPU_OFF && val > TPU_OFF) {
+		ret = pm_runtime_get_sync(dev);
+		if (ret) {
+			dev_err(dev, "%s: pm_runtime_get_sync err: %d\n",
+				__func__, ret);
+			return ret;
+		}
+	}
+
+	/* TPU_OFF is invalid state */
+	if (val != TPU_OFF) {
+		ret = exynos_acpm_set_rate(TPU_ACPM_DOMAIN, (unsigned long)val);
+		if (ret) {
+			dev_err(dev, "error setting tpu power state: %d\n", ret);
+			pm_runtime_put_sync(dev);
+			return ret;
+		}
+	}
+
+	if (curr_state != TPU_OFF && val == TPU_OFF) {
+		ret = pm_runtime_put_sync(dev);
+		if (ret) {
+			dev_err(dev, "%s: pm_runtime_put_sync returned %d\n",
+				__func__, ret);
+			return ret;
+		}
+	}
+
+	return ret;
+}
+
+static int janeiro_pwr_state_get_locked(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct device *dev = etdev->dev;
+
+	*val = exynos_acpm_get_rate(TPU_ACPM_DOMAIN, 0);
+	dev_dbg(dev, "current tpu power state: %llu\n", *val);
+
+	return 0;
+}
+
+static int janeiro_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->requested_state = val;
+	if (val >= platform_pwr->min_state)
+		ret = janeiro_pwr_state_set_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int janeiro_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret;
+
+	mutex_lock(&platform_pwr->state_lock);
+	ret = janeiro_pwr_state_get_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int janeiro_min_pwr_state_set(void *data, u64 val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+	int ret = 0;
+
+	mutex_lock(&platform_pwr->state_lock);
+	platform_pwr->min_state = val;
+	if (val >= platform_pwr->requested_state)
+		ret = janeiro_pwr_state_set_locked(etdev, val);
+	mutex_unlock(&platform_pwr->state_lock);
+	return ret;
+}
+
+static int janeiro_min_pwr_state_get(void *data, u64 *val)
+{
+	struct edgetpu_dev *etdev = (typeof(etdev))data;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	struct janeiro_platform_pwr *platform_pwr = &edgetpu_pdev->platform_pwr;
+
+	mutex_lock(&platform_pwr->state_lock);
+	*val = platform_pwr->min_state;
+	mutex_unlock(&platform_pwr->state_lock);
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, janeiro_pwr_state_get,
+			 janeiro_pwr_state_set, "%llu\n");
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, janeiro_min_pwr_state_get,
+			 janeiro_min_pwr_state_set, "%llu\n");
+
+static int janeiro_get_initial_pwr_state(struct device *dev)
+{
+	switch (power_state) {
+	case TPU_ACTIVE_UUD:
+	case TPU_ACTIVE_SUD:
+	case TPU_ACTIVE_UD:
+	case TPU_ACTIVE_NOM:
+		dev_info(dev, "Initial power state: %d\n", power_state);
+		break;
+	case TPU_OFF:
+		dev_warn(dev, "Power state %d prevents control core booting",
+			 power_state);
+#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 4, 0)
+		fallthrough;
+#endif
+	default:
+		dev_warn(dev, "Power state %d is invalid\n", power_state);
+		dev_warn(dev, "defaulting to active nominal\n");
+		power_state = TPU_ACTIVE_NOM;
+		break;
+	}
+	return power_state;
+}
+
+static void janeiro_power_down(struct edgetpu_pm *etpm);
+
+#define EDGETPU_PSM0_CFG 0x1c1880
+#define EDGETPU_PSM0_START 0x1c1884
+#define EDGETPU_PSM0_STATUS 0x1c1888
+#define EDGETPU_PSM1_CFG 0x1c2880
+#define EDGETPU_PSM1_START 0x1c2884
+#define EDGETPU_PSM1_STATUS 0x1c2888
+#define EDGETPU_LPM_CHANGE_TIMEOUT 30000
+
+static int janeiro_set_lpm(struct edgetpu_dev *etdev)
+{
+	int ret;
+	u32 val;
+
+	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_START, 1);
+	ret = readl_poll_timeout(etdev->regs.mem + EDGETPU_PSM0_STATUS, val,
+				 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, 5, EDGETPU_LPM_CHANGE_TIMEOUT);
+	if (ret) {
+		etdev_err(etdev, "Set LPM1 failed: %d\n", ret);
+		return ret;
+	}
+
+	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM0_CFG, 0);
+	edgetpu_dev_write_32_sync(etdev, EDGETPU_PSM1_CFG, 0);
+
+	return 0;
+}
+
+static int janeiro_power_up(struct edgetpu_pm *etpm)
+{
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	int ret = 0;
+
+	ret = janeiro_pwr_state_set(
+		etpm->etdev, janeiro_get_initial_pwr_state(etdev->dev));
+
+	etdev_info(etpm->etdev, "Powering up\n");
+
+	if (ret)
+		return ret;
+
+	janeiro_set_lpm(etdev);
+
+	/* Clear out log / trace buffers */
+	memset(edgetpu_pdev->log_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
+#if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
+	memset(edgetpu_pdev->trace_mem.vaddr, 0, EDGETPU_TELEMETRY_BUFFER_SIZE);
+#endif
+
+	edgetpu_chip_init(etdev);
+
+	if (etdev->kci) {
+		etdev_dbg(etdev, "Resetting KCI\n");
+		edgetpu_kci_reinit(etdev->kci);
+	}
+	if (etdev->mailbox_manager) {
+		etdev_dbg(etdev, "Resetting VII mailboxes\n");
+		edgetpu_mailbox_reset_vii(etdev->mailbox_manager);
+	}
+
+	if (!etdev->firmware)
+		return 0;
+
+	/*
+	 * Why this function uses edgetpu_firmware_*_locked functions without explicitly holding
+	 * edgetpu_firmware_lock:
+	 *
+	 * edgetpu_pm_get() is called in two scenarios - one is when the firmware loading is
+	 * attempting, another one is when the user-space clients need the device be powered
+	 * (usually through acquiring the wakelock).
+	 *
+	 * For the first scenario edgetpu_firmware_is_loading() below shall return true.
+	 * For the second scenario we are indeed called without holding the firmware lock, but the
+	 * firmware loading procedures (i.e. the first scenario) always call edgetpu_pm_get() before
+	 * changing the firmware state, and edgetpu_pm_get() is blocked until this function
+	 * finishes. In short, we are protected by the PM lock.
+	 */
+
+	if (edgetpu_firmware_is_loading(etdev))
+		return 0;
+
+	/* attempt firmware run */
+	switch (edgetpu_firmware_status_locked(etdev)) {
+	case FW_VALID:
+		ret = edgetpu_firmware_restart_locked(etdev);
+		break;
+	case FW_INVALID:
+		ret = edgetpu_firmware_run_locked(etdev->firmware,
+						  EDGETPU_DEFAULT_FIRMWARE_NAME,
+						  FW_DEFAULT);
+		break;
+	default:
+		break;
+	}
+
+	if (ret)
+		janeiro_power_down(etpm);
+
+	return ret;
+}
+
+static void
+janeiro_pm_shutdown_firmware(struct janeiro_platform_dev *etpdev,
+			     struct edgetpu_dev *etdev)
+{
+	int ret;
+
+	ret = edgetpu_kci_shutdown(etdev->kci);
+	if (ret) {
+		etdev_err(etdev, "firmware shutdown failed: %d",
+			  ret);
+		return;
+	}
+}
+
+static void janeiro_power_down(struct edgetpu_pm *etpm)
+{
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	u64 val;
+
+	etdev_info(etdev, "Powering down\n");
+
+	if (janeiro_pwr_state_get(etdev, &val)) {
+		etdev_warn(etdev, "Failed to read current power state\n");
+		val = TPU_ACTIVE_NOM;
+	}
+	if (val == TPU_OFF) {
+		etdev_dbg(etdev, "Device already off, skipping shutdown\n");
+		return;
+	}
+
+	if (etdev->kci && edgetpu_firmware_status_locked(etdev) == FW_VALID) {
+		/* Update usage stats before we power off fw. */
+		edgetpu_kci_update_usage_locked(etdev);
+		janeiro_pm_shutdown_firmware(edgetpu_pdev, etdev);
+		edgetpu_kci_cancel_work_queues(etdev->kci);
+	}
+
+	janeiro_pwr_state_set(etdev, TPU_OFF);
+}
+
 static int janeiro_pm_after_create(struct edgetpu_pm *etpm)
 {
-	struct device *dev = etpm->etdev->dev;
+	int ret;
+	struct edgetpu_dev *etdev = etpm->etdev;
+	struct janeiro_platform_dev *edgetpu_pdev = to_janeiro_dev(etdev);
+	struct device *dev = etdev->dev;
+
+	ret = janeiro_pwr_state_init(dev);
+	if (ret)
+		return ret;
 
-	return janeiro_pwr_state_init(dev);
+	mutex_init(&edgetpu_pdev->platform_pwr.state_lock);
+
+	ret = janeiro_pwr_state_set(etdev,
+				     janeiro_get_initial_pwr_state(dev));
+	if (ret)
+		return ret;
+	janeiro_pwr_debugfs_dir =
+		debugfs_create_dir("power", edgetpu_fs_debugfs_dir());
+	if (IS_ERR_OR_NULL(janeiro_pwr_debugfs_dir)) {
+		etdev_warn(etdev, "Failed to create debug FS power");
+		/* don't fail the procedure on debug FS creation fails */
+		return 0;
+	}
+	debugfs_create_file("state", 0660, janeiro_pwr_debugfs_dir, etdev,
+			    &fops_tpu_pwr_state);
+	debugfs_create_file("min_state", 0660, janeiro_pwr_debugfs_dir, etdev,
+			    &fops_tpu_min_pwr_state);
+	return 0;
 }
 
 static void janeiro_pm_before_destroy(struct edgetpu_pm *etpm)
 {
+	debugfs_remove_recursive(janeiro_pwr_debugfs_dir);
 	pm_runtime_disable(etpm->etdev->dev);
 }
 
 static struct edgetpu_pm_handlers janeiro_pm_handlers = {
 	.after_create = janeiro_pm_after_create,
 	.before_destroy = janeiro_pm_before_destroy,
+	.power_up = janeiro_power_up,
+	.power_down = janeiro_power_down,
 };
 
 int janeiro_pm_create(struct edgetpu_dev *etdev)
diff --git a/drivers/edgetpu/janeiro-pm.h b/drivers/edgetpu/janeiro-pm.h
index 3a45fde..78991cc 100644
--- a/drivers/edgetpu/janeiro-pm.h
+++ b/drivers/edgetpu/janeiro-pm.h
@@ -7,8 +7,73 @@
 #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.h b/drivers/edgetpu/janeiro/config.h
index 23950fb..0ed255b 100644
--- a/drivers/edgetpu/janeiro/config.h
+++ b/drivers/edgetpu/janeiro/config.h
@@ -40,11 +40,11 @@
 
 /*
  * Instruction remap registers make carveout memory appear at address
- * 0x10000000 from the R52 perspective
+ * 0x10000000 from the TPU CPU perspective
  */
 #define EDGETPU_INSTRUCTION_REMAP_BASE		0x10000000
 
-/* Address from which the R52 can access data in the remapped region */
+/* Address from which the TPU CPU can access data in the remapped region */
 #define EDGETPU_REMAPPED_DATA_ADDR                                        \
 	(EDGETPU_INSTRUCTION_REMAP_BASE + EDGETPU_REMAPPED_DATA_OFFSET)
 #include "config-mailbox.h"