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

  gcip: implement gcip image config utilities
  Bug: 243500340
  gcip: add gcip-image-config.h
  Bug: 243500340 (repeat)
  edgetpu: Program SSMT according to EDGETPU_NUM_SSMTS
  Bug: 241495826
  gcip: Add linux/workqueue.h as header
  edgetpu: move rio SSMT init to new SoC init call for gsx01
  edgetpu: move PM init before MMU attach
  edgetpu: Setting up SSMT to feed-through mode
  Bug: 241495826 (repeat)
  edgetpu: Use gcip.o in Kbuild
  Bug: 240234334
  edgetpu: remove unneeded header include
  edgetpu: rename mobile PM functions to be edgetpu-specific
  edgetpu: move chip PM init call inside device add function
  edgetpu: move gsx01 rkci defines to new header, add gsx01 unit tests
  edgetpu: gsx01: fixup style a.o. issues
  edgetpu: move gsx01 pm_qos and bts functions to gsx01 soc source
  edgetpu: move gsx01-specific pm debugfs interfaces to gsx01 soc source
  edgetpu: support MMU_COHERENT on buffer mapping
  Bug: 241495826 (repeat)
  edgetpu: add future chip and SoC skeletal support
  Bug: 235693240

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

23 files changed, 869 insertions, 496 deletions

diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 292d42a..85e400b 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -13,9 +13,7 @@ endif
 
 GCIP_DIR=gcip-kernel-driver/drivers/gcip
 
-gcip-objs	:= $(GCIP_DIR)/gcip-firmware.o $(GCIP_DIR)/gcip-kci.o $(GCIP_DIR)/gcip-mailbox.o $(GCIP_DIR)/gcip-mem-pool.o
-
-edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-domain-pool.o $(gcip-objs)
+edgetpu-objs	:= edgetpu-mailbox.o edgetpu-kci.o edgetpu-telemetry.o edgetpu-mapping.o edgetpu-dmabuf.o edgetpu-async.o edgetpu-iremap-pool.o edgetpu-sw-watchdog.o edgetpu-firmware.o edgetpu-firmware-util.o edgetpu-domain-pool.o $(GCIP_DIR)/gcip.o
 
 rio-y		:= rio-device.o rio-device-group.o rio-fs.o rio-core.o rio-platform.o rio-firmware.o rio-thermal.o rio-pm.o rio-debug-dump.o rio-usage-stats.o rio-iommu.o rio-wakelock.o rio-external.o rio-soc.o $(edgetpu-objs)
 
@@ -33,3 +31,5 @@ CFLAGS_rio-usage-stats.o := -DCONFIG_RIO=1
 CFLAGS_rio-wakelock.o := -DCONFIG_RIO=1
 CFLAGS_rio-external.o := -DCONFIG_RIO=1
 CFLAGS_rio-soc.o := -DCONFIG_RIO=1
+
+
diff --git a/drivers/edgetpu/Makefile b/drivers/edgetpu/Makefile
index acd14bd..27ca3a1 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -27,10 +27,19 @@ rio-objs	:= rio-core.o rio-debug-dump.o rio-device-group.o rio-device.o \
 		   rio-thermal.o rio-usage-stats.o rio-wakelock.o rio-external.o \
 		   rio-soc.o $(edgetpu-objs)
 
+
 ccflags-y	+= -I$(KERNEL_SRC)/../private/google-modules/edgetpu/rio/drivers/edgetpu/include
 KBUILD_OPTIONS	+= CONFIG_RIO=m
 include $(KERNEL_SRC)/../private/google-modules/soc/gs/Makefile.include
 
-modules modules_install clean:
+GCIP_DIR=gcip-kernel-driver/drivers/gcip
+
+modules modules_install:
+	$(MAKE) -C $(KERNEL_SRC) M=$(M)/$(GCIP_DIR) W=1 gcip.o
+	$(MAKE) -C $(KERNEL_SRC) M=$(M) W=1 $(KBUILD_OPTIONS) \
+	EXTRA_CFLAGS="$(EXTRA_CFLAGS)" KBUILD_EXTRA_SYMBOLS="$(EXTRA_SYMBOLS)" $(@)
+clean:
+	$(MAKE) -C $(KERNEL_SRC) M=$(M)/$(GCIP_DIR) W=1 \
+	EXTRA_CFLAGS="$(EXTRA_CFLAGS)" KBUILD_EXTRA_SYMBOLS="$(EXTRA_SYMBOLS)" $(@)
 	$(MAKE) -C $(KERNEL_SRC) M=$(M) W=1 $(KBUILD_OPTIONS) \
 	EXTRA_CFLAGS="$(EXTRA_CFLAGS)" KBUILD_EXTRA_SYMBOLS="$(EXTRA_SYMBOLS)" $(@)
diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index 9aa87bb..c496a33 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -25,6 +25,10 @@
 #define EDGETPU_NUM_CORES 1
 #endif
 
+#ifndef EDGETPU_NUM_SSMTS
+#define EDGETPU_NUM_SSMTS 1
+#endif
+
 /* Uses a smaller size for unittests to avoid DMA warnings. */
 #if IS_ENABLED(CONFIG_EDGETPU_TEST)
 #undef EDGETPU_DEBUG_DUMP_MEM_SIZE
diff --git a/drivers/edgetpu/edgetpu-core.c b/drivers/edgetpu/edgetpu-core.c
index 0aa4fd0..dc8ba6c 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -446,9 +446,17 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 			etdev->dev_name, ret);
 		goto remove_dev;
 	}
+
+	/* Init PM in case the platform needs power up actions before MMU setup and such. */
+	ret = edgetpu_chip_pm_create(etdev);
+	if (ret) {
+		etdev_err(etdev, "Failed to initialize PM interface: %d", ret);
+		goto remove_mboxes;
+	}
+
 	ret = edgetpu_chip_setup_mmu(etdev);
 	if (ret)
-		goto remove_dev;
+		goto remove_mboxes;
 
 	edgetpu_usage_stats_init(etdev);
 
@@ -481,6 +489,8 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 remove_usage_stats:
 	edgetpu_usage_stats_exit(etdev);
 	edgetpu_chip_remove_mmu(etdev);
+remove_mboxes:
+	edgetpu_mailbox_remove_all(etdev->mailbox_manager);
 remove_dev:
 	edgetpu_mark_probe_fail(etdev);
 	edgetpu_fs_remove(etdev);
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index a8f6423..a9260e7 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -170,6 +170,7 @@ struct edgetpu_dev {
 	struct device *dev;	   /* platform/pci bus device */
 	uint num_ifaces;		   /* Number of device interfaces */
 	uint num_cores; /* Number of cores */
+	uint num_ssmts; /* Number of SSMTs */
 	/*
 	 * Array of device interfaces
 	 * First element is the default interface
@@ -405,14 +406,6 @@ int edgetpu_chip_setup_mmu(struct edgetpu_dev *etdev);
  */
 void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev);
 
-/*
- * Handle chip-specific incoming requests from firmware over KCI
- * Note: This will get called from the system's work queue.
- * Code should not block for extended periods of time
- */
-void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
-				     struct gcip_kci_response_element *resp);
-
 /* Device -> Core API */
 
 /* Add current thread as new TPU client */
diff --git a/drivers/edgetpu/edgetpu-kci.c b/drivers/edgetpu/edgetpu-kci.c
index 18dc690..16e4a97 100644
--- a/drivers/edgetpu/edgetpu-kci.c
+++ b/drivers/edgetpu/edgetpu-kci.c
@@ -20,6 +20,7 @@
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mmu.h"
+#include "edgetpu-soc.h"
 #include "edgetpu-telemetry.h"
 #include "edgetpu-usage-stats.h"
 
@@ -145,7 +146,7 @@ static void edgetpu_reverse_kci_handle_response(struct gcip_kci *kci,
 	struct edgetpu_dev *etdev = mailbox->etdev;
 
 	if (resp->code <= GCIP_RKCI_CHIP_CODE_LAST) {
-		edgetpu_chip_handle_reverse_kci(etdev, resp);
+		edgetpu_soc_handle_reverse_kci(etdev, resp);
 		return;
 	}
 
diff --git a/drivers/edgetpu/edgetpu-mapping.h b/drivers/edgetpu/edgetpu-mapping.h
index ae42a3a..65547bc 100644
--- a/drivers/edgetpu/edgetpu-mapping.h
+++ b/drivers/edgetpu/edgetpu-mapping.h
@@ -160,11 +160,8 @@ static inline int mmu_flag_to_iommu_prot(u32 mmu_flags, struct device *dev,
 {
 	int prot = 0;
 
-	if (mmu_flags & EDGETPU_MMU_COHERENT) {
-#ifdef EDGETPU_IS_DMA_COHERENT
+	if (mmu_flags & EDGETPU_MMU_COHERENT)
 		prot = IOMMU_CACHE;
-#endif
-	}
 	prot |= __dma_dir_to_iommu_prot(dir);
 	return prot;
 }
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index 735afdb..ce7e75c 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -18,6 +18,7 @@
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-mmu.h"
 #include "edgetpu-mobile-platform.h"
+#include "edgetpu-soc.h"
 #include "edgetpu-telemetry.h"
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
@@ -220,29 +221,6 @@ void edgetpu_chip_remove_mmu(struct edgetpu_dev *etdev)
 	edgetpu_mmu_detach(etdev);
 }
 
-static int edgetpu_platform_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
-{
-	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
-	struct platform_device *pdev = to_platform_device(etdev->dev);
-	struct resource *res;
-	int ret;
-	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)) {
-		ret = PTR_ERR(ssmt_base);
-		etdev_warn(etdev, "Failed to map SSMT register base: %d", ret);
-		return ret;
-	}
-	etmdev->ssmt_base = ssmt_base;
-	return 0;
-}
-
 static int edgetpu_platform_setup_irq(struct edgetpu_mobile_platform_dev *etmdev)
 {
 	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
@@ -331,6 +309,7 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 	platform_set_drvdata(pdev, etdev);
 	etdev->dev = dev;
 	etdev->num_cores = EDGETPU_NUM_CORES;
+	etdev->num_ssmts = EDGETPU_NUM_SSMTS;
 
 	r = platform_get_resource(pdev, IORESOURCE_MEM, 0);
 	if (IS_ERR_OR_NULL(r)) {
@@ -346,15 +325,11 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 		return -ENODEV;
 	}
 
+	edgetpu_soc_init(etdev);
+
 	mutex_init(&etmdev->platform_pwr.policy_lock);
 	etmdev->platform_pwr.curr_policy = TPU_POLICY_MAX;
 
-	ret = edgetpu_chip_pm_create(etdev);
-	if (ret) {
-		dev_err(dev, "Failed to initialize PM interface: %d", ret);
-		return ret;
-	}
-
 	ret = edgetpu_platform_setup_fw_region(etmdev);
 	if (ret) {
 		dev_err(dev, "setup fw regions failed: %d", ret);
@@ -391,10 +366,6 @@ static int edgetpu_mobile_platform_probe(struct platform_device *pdev,
 		goto out_remove_device;
 	}
 
-	ret = edgetpu_platform_parse_ssmt(etmdev);
-	if (ret)
-		dev_warn(dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
-
 	etmdev->log_mem = devm_kcalloc(dev, etdev->num_cores, sizeof(*etmdev->log_mem), GFP_KERNEL);
 	if (!etmdev->log_mem) {
 		ret = -ENOMEM;
@@ -469,6 +440,6 @@ static int edgetpu_mobile_platform_remove(struct platform_device *pdev)
 	edgetpu_platform_cleanup_fw_region(etmdev);
 	edgetpu_pm_put(etdev->pm);
 	edgetpu_pm_shutdown(etdev, true);
-	mobile_pm_destroy(etdev);
+	edgetpu_mobile_pm_destroy(etdev);
 	return 0;
 }
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index 3aa0104..afdaa9d 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -80,7 +80,7 @@ struct edgetpu_mobile_platform_dev {
 	 */
 	struct device *gsa_dev;
 	/* Virtual address of the SSMT block for this chip. */
-	void __iomem *ssmt_base;
+	void __iomem **ssmt_base;
 	/* Coherent log buffer */
 	struct edgetpu_coherent_mem *log_mem;
 	/* Coherent trace buffer */
diff --git a/drivers/edgetpu/edgetpu-pm.h b/drivers/edgetpu/edgetpu-pm.h
index 96e4a4d..4acdb45 100644
--- a/drivers/edgetpu/edgetpu-pm.h
+++ b/drivers/edgetpu/edgetpu-pm.h
@@ -72,6 +72,12 @@ void edgetpu_pm_put(struct edgetpu_pm *etpm);
 int edgetpu_pm_create(struct edgetpu_dev *etdev,
 		      const struct edgetpu_pm_handlers *handlers);
 
+/*
+ * Wrapper for chip-specific implementation.
+ * Typically calls mobile_pm_create after initializing the platform_pwr struct.
+ */
+int edgetpu_chip_pm_create(struct edgetpu_dev *etdev);
+
 /* Destroy the power management interface associated with an edgetpu device */
 void edgetpu_pm_destroy(struct edgetpu_dev *etdev);
 
diff --git a/drivers/edgetpu/edgetpu-soc.h b/drivers/edgetpu/edgetpu-soc.h
index 31083a4..5e87676 100644
--- a/drivers/edgetpu/edgetpu-soc.h
+++ b/drivers/edgetpu/edgetpu-soc.h
@@ -11,10 +11,14 @@
 #include <linux/types.h>
 
 #include "edgetpu-internal.h"
+#include "edgetpu-kci.h"
 #include "edgetpu-thermal.h"
 
 /* SoC-specific calls for the following functions. */
 
+/* Probe-time init */
+void edgetpu_soc_init(struct edgetpu_dev *etdev);
+
 /* Prep for running firmware: set access control, etc. */
 int edgetpu_soc_prepare_firmware(struct edgetpu_dev *etdev);
 
@@ -30,6 +34,23 @@ int edgetpu_soc_pm_set_init_freq(unsigned long freq);
 /* Set PM policy */
 int edgetpu_soc_pm_set_policy(u64 val);
 
+/* Power down */
+void edgetpu_soc_pm_power_down(struct edgetpu_dev *etdev);
+
+/* Init SoC PM system */
+int edgetpu_soc_pm_init(struct edgetpu_dev *etdev);
+
+/* De-init SoC PM system */
+void edgetpu_soc_pm_exit(struct edgetpu_dev *etdev);
+
+/*
+ * Handle Reverse KCI commands for SoC family.
+ * Note: This will get called from the system's work queue.
+ * Code should not block for extended periods of time
+ */
+void edgetpu_soc_handle_reverse_kci(struct edgetpu_dev *etdev,
+				    struct gcip_kci_response_element *resp);
+
 /* Init thermal subsystem SoC specifics for TPU */
 void edgetpu_soc_thermal_init(struct edgetpu_thermal *thermal);
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index 3056519..1b988e4 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -6,7 +6,7 @@
 CONFIG_GCIP		?= m
 obj-$(CONFIG_GCIP)	+= gcip.o
 
-gcip-objs		:= gcip-firmware.o gcip-kci.o gcip-mailbox.o gcip-mem-pool.o
+gcip-objs		:= gcip-firmware.o gcip-image-config.o gcip-kci.o gcip-mailbox.o gcip-mem-pool.o
 
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
new file mode 100644
index 0000000..af082b6
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
@@ -0,0 +1,206 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Framework for parsing the firmware image configuration.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#include <linux/device.h>
+#include <linux/errno.h>
+#include <linux/string.h>
+#include <linux/types.h>
+
+#include <gcip/gcip-image-config.h>
+
+#define SIZE_MASK 0xfff
+
+/* used by iommu_mappings */
+#define CONFIG_TO_SIZE(a) ((1U << ((a) & SIZE_MASK)) << 12)
+
+/* used by ns_iommu_mappings */
+#define CONFIG_TO_MBSIZE(a) (((a) & SIZE_MASK) << 20)
+
+static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
+				struct gcip_image_config *config)
+{
+	int i, ret;
+	dma_addr_t daddr;
+	size_t size;
+	phys_addr_t paddr;
+
+	for (i = 0; i < config->num_iommu_mappings; i++) {
+		daddr = config->iommu_mappings[i].virt_address;
+		if (unlikely(!daddr)) {
+			dev_warn(parser->dev, "Invalid config, device address is zero");
+			ret = -EIO;
+			goto err;
+		}
+		size = CONFIG_TO_SIZE(config->iommu_mappings[i].image_config_value);
+		paddr = config->iommu_mappings[i].image_config_value & ~SIZE_MASK;
+
+		dev_dbg(parser->dev, "Image config adding IOMMU mapping: %#llx -> %#llx", daddr,
+			paddr);
+
+		if (unlikely(daddr + size <= daddr || paddr + size <= paddr)) {
+			ret = -EOVERFLOW;
+			goto err;
+		}
+		ret = parser->ops->map(parser->data, daddr, paddr, size,
+				       GCIP_IMAGE_CONFIG_FLAGS_SECURE);
+		if (ret) {
+			dev_err(parser->dev,
+				"Unable to Map: %d dma_addr: %#llx phys_addr: %#llx size: %#lx\n",
+				ret, daddr, paddr, size);
+			goto err;
+		}
+	}
+
+	return 0;
+
+err:
+	while (i--) {
+		daddr = config->iommu_mappings[i].virt_address;
+		size = CONFIG_TO_SIZE(config->iommu_mappings[i].image_config_value);
+		parser->ops->unmap(parser->data, daddr, size, GCIP_IMAGE_CONFIG_FLAGS_SECURE);
+	}
+	return ret;
+}
+
+static void clear_iommu_mappings(struct gcip_image_config_parser *parser,
+				 struct gcip_image_config *config)
+{
+	dma_addr_t daddr;
+	size_t size;
+	int i;
+
+	for (i = config->num_iommu_mappings - 1; i >= 0; i--) {
+		daddr = config->iommu_mappings[i].virt_address;
+		size = CONFIG_TO_SIZE(config->iommu_mappings[i].image_config_value);
+		dev_dbg(parser->dev, "Image config removing IOMMU mapping: %#llx size=%#lx", daddr,
+			size);
+		parser->ops->unmap(parser->data, daddr, size, GCIP_IMAGE_CONFIG_FLAGS_SECURE);
+	}
+}
+
+static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
+				   struct gcip_image_config *config)
+{
+	dma_addr_t daddr;
+	size_t size;
+	int ret, i;
+	phys_addr_t paddr = 0;
+
+	for (i = 0; i < config->num_ns_iommu_mappings; i++) {
+		daddr = config->ns_iommu_mappings[i] & ~SIZE_MASK;
+		if (unlikely(!daddr)) {
+			dev_warn(parser->dev, "Invalid config, device address is zero");
+			ret = -EIO;
+			goto err;
+		}
+		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
+		dev_dbg(parser->dev, "Image config adding NS IOMMU mapping: %#llx -> %#llx", daddr,
+			paddr);
+		if (unlikely(daddr + size <= daddr || paddr + size <= paddr)) {
+			ret = -EOVERFLOW;
+			goto err;
+		}
+		ret = parser->ops->map(parser->data, daddr, paddr, size, 0);
+		if (ret)
+			goto err;
+		paddr += size;
+	}
+
+	return 0;
+
+err:
+	while (i--) {
+		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
+		daddr = config->ns_iommu_mappings[i] & ~SIZE_MASK;
+		parser->ops->unmap(parser->data, daddr, size, 0);
+	}
+	return ret;
+}
+
+static void clear_ns_iommu_mappings(struct gcip_image_config_parser *parser,
+				    struct gcip_image_config *config)
+{
+	dma_addr_t daddr;
+	size_t size;
+	int i;
+
+	for (i = config->num_ns_iommu_mappings - 1; i >= 0; i--) {
+		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
+		daddr = config->ns_iommu_mappings[i] & ~SIZE_MASK;
+		dev_dbg(parser->dev, "Image config removing NS IOMMU mapping: %#llx size=%#lx",
+			daddr, size);
+		parser->ops->unmap(parser->data, daddr, size, 0);
+	}
+}
+
+static int map_image_config(struct gcip_image_config_parser *parser,
+			    struct gcip_image_config *config)
+{
+	int ret = setup_ns_iommu_mappings(parser, config);
+
+	if (ret)
+		return ret;
+	if (gcip_image_config_is_ns(config)) {
+		ret = setup_iommu_mappings(parser, config);
+		if (ret)
+			clear_ns_iommu_mappings(parser, config);
+	}
+	return ret;
+}
+
+static void unmap_image_config(struct gcip_image_config_parser *parser,
+			       struct gcip_image_config *config)
+{
+	if (gcip_image_config_is_ns(config))
+		clear_iommu_mappings(parser, config);
+	clear_ns_iommu_mappings(parser, config);
+}
+
+int gcip_image_config_parser_init(struct gcip_image_config_parser *parser,
+				  const struct gcip_image_config_ops *ops, struct device *dev,
+				  void *data)
+{
+	if (!ops->map || !ops->unmap) {
+		dev_err(dev, "Missing mandatory operations for image config parser");
+		return -EINVAL;
+	}
+	parser->dev = dev;
+	parser->data = data;
+	parser->ops = ops;
+	memset(&parser->last_config, 0, sizeof(parser->last_config));
+	return 0;
+}
+
+int gcip_image_config_parse(struct gcip_image_config_parser *parser,
+			    struct gcip_image_config *config)
+{
+	int ret;
+
+	if (!memcmp(config, &parser->last_config, sizeof(*config)))
+		return 0;
+	unmap_image_config(parser, &parser->last_config);
+	ret = map_image_config(parser, config);
+	if (ret) {
+		dev_err(parser->dev, "Map image config failed: %d", ret);
+		/*
+		 * Weird case as the mappings in the last config were just removed - might happen
+		 * if the IOMMU driver state is corrupted. We can't help to rescue it so let's
+		 * simply log a message.
+		 */
+		if (unlikely(map_image_config(parser, &parser->last_config)))
+			dev_err(parser->dev, "Failed to roll back the last image config");
+		return ret;
+	}
+	memcpy(&parser->last_config, config, sizeof(parser->last_config));
+	return 0;
+}
+
+void gcip_image_config_clear(struct gcip_image_config_parser *parser)
+{
+	unmap_image_config(parser, &parser->last_config);
+	memset(&parser->last_config, 0, sizeof(parser->last_config));
+}
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
new file mode 100644
index 0000000..6c3333e
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
@@ -0,0 +1,150 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Framework for parsing the firmware image configuration.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#ifndef __GCIP_IMAGE_CONFIG_H__
+#define __GCIP_IMAGE_CONFIG_H__
+
+#include <linux/types.h>
+
+#define GCIP_FW_NUM_VERSIONS 4
+#define GCIP_IMG_CFG_MAX_IOMMU_MAPPINGS 23
+#define GCIP_IMG_CFG_MAX_NS_IOMMU_MAPPINGS 5
+
+#define GCIP_FW_PRIV_LEVEL_GSA 0
+#define GCIP_FW_PRIV_LEVEL_TZ 1
+#define GCIP_FW_PRIV_LEVEL_NS 2
+
+/*
+ * The image configuration attached to the signed firmware.
+ */
+struct gcip_image_config {
+	__u32 carveout_base;
+	__u32 firmware_base;
+	__u32 firmware_size;
+	__u32 firmware_versions[GCIP_FW_NUM_VERSIONS];
+	__u32 config_version;
+	__u32 privilege_level;
+	__u32 remapped_region_start;
+	__u32 remapped_region_size;
+	__u32 num_iommu_mappings;
+	struct {
+		/* Device virtual address */
+		__u32 virt_address;
+		/*
+		 * contains a 12-bit aligned address and a page-order size into a
+		 * 32-bit value i.e. a physical address and size in page order.
+		 */
+		__u32 image_config_value;
+	} iommu_mappings[GCIP_IMG_CFG_MAX_IOMMU_MAPPINGS];
+	__u32 num_ns_iommu_mappings;
+	__u32 ns_iommu_mappings[GCIP_IMG_CFG_MAX_NS_IOMMU_MAPPINGS];
+} __packed;
+
+#define GCIP_IMAGE_CONFIG_FLAGS_SECURE (1u << 0)
+
+struct gcip_image_config_ops {
+	/*
+	 * Adds an IOMMU mapping from @daddr to @paddr with size @size.
+	 *
+	 * It is ensured that there is no overflow on @paddr + @size before calling this function.
+	 *
+	 * @flags is a bit-field with the following attributes:
+	 *   [0:0] - Security. 1 for secure and 0 for non-secure.
+	 *   [31:1] - Reserved.
+	 *
+	 * Returns 0 on success. Otherwise a negative errno.
+	 * Mandatory.
+	 */
+	int (*map)(void *data, dma_addr_t daddr, phys_addr_t paddr, size_t size,
+		   unsigned int flags);
+	/*
+	 * Removes the IOMMU mapping previously added by @map.
+	 *
+	 * Mandatory.
+	 */
+	void (*unmap)(void *data, dma_addr_t daddr, size_t size, unsigned int flags);
+};
+
+struct gcip_image_config_parser {
+	struct device *dev;
+	void *data; /* User-specify data, will be passed to ops. */
+	const struct gcip_image_config_ops *ops;
+	/* The last image config being successfully parsed. */
+	struct gcip_image_config last_config;
+};
+
+/*
+ * Initializes the image configuration parser.
+ *
+ * @dev is only used for logging.
+ * @data will be passed to operations.
+ *
+ * Returns 0 on success. Returns -EINVAL when any mandatory operations is NULL.
+ */
+int gcip_image_config_parser_init(struct gcip_image_config_parser *parser,
+				  const struct gcip_image_config_ops *ops, struct device *dev,
+				  void *data);
+
+/*
+ * Parses the image configuration and adds specified IOMMU mappings by calling pre-registered
+ * operations.
+ *
+ * Number of mappings to be added might be different according to the value of
+ * @config->privilege_level:
+ * - GCIP_FW_PRIV_LEVEL_NS:
+ *   Both @iommu_mappings and @ns_iommu_mappings will be added. Because GCIP_FW_PRIV_LEVEL_NS means
+ *   the firmware will run in non-secure mode and all transactions will go through the non-secure
+ *   IOMMU.
+ * - Otherwise:
+ *   Only @ns_iommu_mappings are considered. TZ/GSA will be the one who programs secure IOMMU for
+ *   those secure IOMMU mappings.
+ *
+ * Before parsing the newly passed @config, the mappings of the last record (stored by @parser
+ * internally) will be reverted. If there is any mapping in the new config fails to be mapped, the
+ * reverted last config will be reverted again. i.e. This function will keep the mapping state the
+ * same as before calling it on any error happens. But if the IOMMU state is somehow corrupted and
+ * hence fails to roll back the reverted last image config, only an error is logged. See the pseudo
+ * code below:
+ *
+ * gcip_image_config_parse(config):
+ *   unmap(last_image_config)
+ *   if ret = map(config) fails:
+ *     LOG("Failed to map image config, rolling back to the last image config.")
+ *     if map(last_image_config) fails:
+ *       LOG("Failed to roll back the last image config.")
+ *     return ret
+ *  else:
+ *    last_image_config = config
+ *  return SUCCESS
+ *
+ * A special case being considered is if the content of @config is identical to the last
+ * successfully parsed image config, this function will return 0 immediately without removing /
+ * adding any mapping.
+ *
+ * Returns 0 on success. Otherwise an errno, which usually would be the one returned by
+ * gcip_image_config_ops.map. On error no new mapping specified in @config is added.
+ */
+int gcip_image_config_parse(struct gcip_image_config_parser *parser,
+			    struct gcip_image_config *config);
+
+/*
+ * Clears the mappings specified in the last image config.
+ *
+ * It's valid to call this function without any image config has been successfully parsed, or when
+ * the last image config is already cleared. In which case this function works as no-op.
+ */
+void gcip_image_config_clear(struct gcip_image_config_parser *parser);
+
+/*
+ * Returns whether the privilege level specified by @config is non-secure.
+ */
+static inline bool gcip_image_config_is_ns(struct gcip_image_config *config)
+{
+	return config->privilege_level == GCIP_FW_PRIV_LEVEL_NS;
+}
+
+#endif /* __GCIP_IMAGE_CONFIG_H__ */
diff --git a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
index 1f54938..71a1e27 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-mailbox.h
@@ -13,6 +13,7 @@
 #include <linux/spinlock.h>
 #include <linux/types.h>
 #include <linux/wait.h>
+#include <linux/workqueue.h>
 
 #define CIRC_QUEUE_WRAPPED(idx, wrap_bit) ((idx)&wrap_bit)
 #define CIRC_QUEUE_INDEX_MASK(wrap_bit) (wrap_bit - 1)
diff --git a/drivers/edgetpu/include/soc/google/bcl.h b/drivers/edgetpu/include/soc/google/bcl.h
deleted file mode 100644
index 506982c..0000000
--- a/drivers/edgetpu/include/soc/google/bcl.h
+++ /dev/null
@@ -1,37 +0,0 @@
-/* SPDX-License-Identifier: GPL-2.0 */
-
-#ifndef __BCL_H__
-#define __BCL_H__
-
-struct gs101_bcl_dev;
-
-static inline unsigned int gs101_get_mpmm(struct gs101_bcl_dev *data)
-{
-	return 0;
-}
-static inline unsigned int gs101_get_ppm(struct gs101_bcl_dev *data)
-{
-	return 0;
-}
-static inline int gs101_set_ppm(struct gs101_bcl_dev *data, unsigned int value)
-{
-	return 0;
-}
-static inline int gs101_set_mpmm(struct gs101_bcl_dev *data, unsigned int value)
-{
-	return 0;
-}
-static inline struct gs101_bcl_dev *gs101_retrieve_bcl_handle(void)
-{
-	return NULL;
-}
-static inline int gs101_init_gpu_ratio(struct gs101_bcl_dev *data)
-{
-	return 0;
-}
-static inline int gs101_init_tpu_ratio(struct gs101_bcl_dev *data)
-{
-	return 0;
-}
-
-#endif /* __BCL_H__ */
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index 5b6c979..d145c0c 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -11,8 +11,6 @@
 #include <linux/module.h>
 #include <linux/pm_runtime.h>
 #include <soc/google/bcl.h>
-#include <soc/google/bts.h>
-#include <soc/google/exynos_pm_qos.h>
 
 #include "edgetpu-config.h"
 #include "edgetpu-firmware.h"
@@ -27,20 +25,10 @@
 #include "edgetpu-pm.c"
 #include "edgetpu-soc.h"
 
-/*
- * Encode INT/MIF values as a 16 bit pair in the 32-bit return value
- * (in units of MHz, to provide enough range)
- */
-#define PM_QOS_INT_SHIFT                (16)
-#define PM_QOS_MIF_MASK                 (0xFFFF)
-#define PM_QOS_FACTOR                   (1000)
-
 static int power_state = TPU_DEFAULT_POWER_STATE;
 
 module_param(power_state, int, 0660);
 
-#define MAX_VOLTAGE_VAL 1250000
-
 enum edgetpu_pwr_state edgetpu_active_states[EDGETPU_NUM_STATES] = {
 	TPU_ACTIVE_UUD,
 	TPU_ACTIVE_SUD,
@@ -50,11 +38,6 @@ enum edgetpu_pwr_state edgetpu_active_states[EDGETPU_NUM_STATES] = {
 
 uint32_t *edgetpu_states_display = edgetpu_active_states;
 
-/* TODO(b/240363978): Remove once ACPM is ready. */
-#if !IS_ENABLED(CONFIG_EDGETPU_TEST)
-unsigned long exynos_acpm_rate = 0;
-#endif /* IS_ENABLED(CONFIG_EDGETPU_TEST) */
-
 static int mobile_pwr_state_init(struct device *dev)
 {
 	int ret;
@@ -83,139 +66,6 @@ static int mobile_pwr_state_init(struct device *dev)
 	return ret;
 }
 
-static int edgetpu_core_rate_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG);
-	return 0;
-}
-
-static int edgetpu_core_rate_set(void *data, u64 val)
-{
-	unsigned long dbg_rate_req;
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG;
-	dbg_rate_req |= val;
-
-	return edgetpu_soc_pm_set_rate(dbg_rate_req);
-}
-
-static int edgetpu_ctl_rate_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG);
-	return 0;
-}
-
-static int edgetpu_ctl_rate_set(void *data, u64 val)
-{
-	unsigned long dbg_rate_req;
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG;
-	dbg_rate_req |= 1000;
-
-	return edgetpu_soc_pm_set_rate(dbg_rate_req);
-}
-
-static int edgetpu_axi_rate_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG);
-	return 0;
-}
-
-static int edgetpu_axi_rate_set(void *data, u64 val)
-{
-	unsigned long dbg_rate_req;
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG;
-	dbg_rate_req |= 1000;
-
-	return edgetpu_soc_pm_set_rate(dbg_rate_req);
-}
-
-static int edgetpu_apb_rate_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_APB_DEBUG);
-	return 0;
-}
-
-static int edgetpu_uart_rate_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_UART_DEBUG);
-	return 0;
-}
-
-static int edgetpu_vdd_int_m_set(void *data, u64 val)
-{
-	struct device *dev = (struct device *)data;
-	unsigned long dbg_rate_req;
-
-	if (val > MAX_VOLTAGE_VAL) {
-		dev_err(dev, "Preventing INT_M voltage > %duV",
-			MAX_VOLTAGE_VAL);
-		return -EINVAL;
-	}
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG;
-	dbg_rate_req |= val;
-
-	return edgetpu_soc_pm_set_rate(dbg_rate_req);
-}
-
-static int edgetpu_vdd_int_m_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG);
-	return 0;
-}
-
-static int edgetpu_vdd_tpu_set(void *data, u64 val)
-{
-	int ret;
-	struct device *dev = (struct device *)data;
-	unsigned long dbg_rate_req;
-
-	if (val > MAX_VOLTAGE_VAL) {
-		dev_err(dev, "Preventing VDD_TPU voltage > %duV",
-			MAX_VOLTAGE_VAL);
-		return -EINVAL;
-	}
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG;
-	dbg_rate_req |= val;
-
-	ret = edgetpu_soc_pm_set_rate(dbg_rate_req);
-	return ret;
-}
-
-static int edgetpu_vdd_tpu_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG);
-	return 0;
-}
-
-static int edgetpu_vdd_tpu_m_set(void *data, u64 val)
-{
-	int ret;
-	struct device *dev = (struct device *)data;
-	unsigned long dbg_rate_req;
-
-	if (val > MAX_VOLTAGE_VAL) {
-		dev_err(dev, "Preventing VDD_TPU voltage > %duV",
-			MAX_VOLTAGE_VAL);
-		return -EINVAL;
-	}
-
-	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG;
-	dbg_rate_req |= val;
-
-	ret = edgetpu_soc_pm_set_rate(dbg_rate_req);
-	return ret;
-}
-
-static int edgetpu_vdd_tpu_m_get(void *data, u64 *val)
-{
-	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG);
-	return 0;
-}
-
 static int mobile_pwr_state_set_locked(struct edgetpu_mobile_platform_dev *etmdev, u64 val)
 {
 	int ret;
@@ -366,29 +216,6 @@ DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_pwr_state, mobile_pwr_state_get, mobile_pwr_st
 DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_min_pwr_state, mobile_min_pwr_state_get, mobile_min_pwr_state_set,
 			"%llu\n");
 
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_core_rate, edgetpu_core_rate_get,
-			 edgetpu_core_rate_set, "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_ctl_rate, edgetpu_ctl_rate_get,
-			 edgetpu_ctl_rate_set, "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_axi_rate, edgetpu_axi_rate_get,
-			 edgetpu_axi_rate_set, "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_apb_rate, edgetpu_apb_rate_get, NULL,
-			 "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_uart_rate, edgetpu_uart_rate_get, NULL,
-			 "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_int_m, edgetpu_vdd_int_m_get,
-			 edgetpu_vdd_int_m_set, "%llu\n");
-
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu, edgetpu_vdd_tpu_get,
-			 edgetpu_vdd_tpu_set, "%llu\n");
-DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu_m, edgetpu_vdd_tpu_m_get,
-			 edgetpu_vdd_tpu_m_set, "%llu\n");
-
 static int mobile_get_initial_pwr_state(struct device *dev)
 {
 	switch (power_state) {
@@ -484,30 +311,6 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 	return ret;
 }
 
-static void mobile_pm_cleanup_bts_scenario(struct edgetpu_dev *etdev)
-{
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
-	int performance_scenario = platform_pwr->performance_scenario;
-
-	if (!performance_scenario)
-		return;
-
-	mutex_lock(&platform_pwr->scenario_lock);
-	while (platform_pwr->scenario_count) {
-		int ret = bts_del_scenario(performance_scenario);
-
-		if (ret) {
-			platform_pwr->scenario_count = 0;
-			etdev_warn_once(etdev, "error %d in cleaning up BTS scenario %u\n", ret,
-					performance_scenario);
-			break;
-		}
-		platform_pwr->scenario_count--;
-	}
-	mutex_unlock(&platform_pwr->scenario_lock);
-}
-
 static void mobile_firmware_down(struct edgetpu_dev *etdev)
 {
 	int ret = edgetpu_kci_shutdown(etdev->etkci);
@@ -569,11 +372,7 @@ static int mobile_power_down(struct edgetpu_pm *etpm)
 
 	mobile_pwr_state_set(etdev, TPU_OFF);
 
-	/* Remove our vote for INT/MIF state (if any) */
-	exynos_pm_qos_update_request(&platform_pwr->int_min, 0);
-	exynos_pm_qos_update_request(&platform_pwr->mif_min, 0);
-
-	mobile_pm_cleanup_bts_scenario(etdev);
+	edgetpu_soc_pm_power_down(etdev);
 
 	/*
 	 * It should be impossible that power_down() is called when secure_client is set.
@@ -602,14 +401,6 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 	mutex_init(&platform_pwr->state_lock);
 	mutex_init(&platform_pwr->scenario_lock);
 
-	exynos_pm_qos_add_request(&platform_pwr->int_min, PM_QOS_DEVICE_THROUGHPUT, 0);
-	exynos_pm_qos_add_request(&platform_pwr->mif_min, PM_QOS_BUS_THROUGHPUT, 0);
-
-	platform_pwr->performance_scenario = bts_get_scenindex("tpu_performance");
-	if (!platform_pwr->performance_scenario)
-		dev_warn(etdev->dev, "tpu_performance BTS scenario not found\n");
-	platform_pwr->scenario_count = 0;
-
 	ret = mobile_pwr_state_set(etdev, mobile_get_initial_pwr_state(dev));
 	if (ret)
 		return ret;
@@ -623,14 +414,7 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 	debugfs_create_file("min_state", 0660, platform_pwr->debugfs_dir, etdev,
 			    &fops_tpu_min_pwr_state);
 	debugfs_create_file("policy", 0660, platform_pwr->debugfs_dir, etdev, &fops_tpu_pwr_policy);
-	debugfs_create_file("vdd_tpu", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_tpu);
-	debugfs_create_file("vdd_tpu_m", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_tpu_m);
-	debugfs_create_file("vdd_int_m", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_int_m);
-	debugfs_create_file("core_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_core_rate);
-	debugfs_create_file("ctl_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_ctl_rate);
-	debugfs_create_file("axi_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_axi_rate);
-	debugfs_create_file("apb_rate", 0440, platform_pwr->debugfs_dir, dev, &fops_tpu_apb_rate);
-	debugfs_create_file("uart_rate", 0440, platform_pwr->debugfs_dir, dev, &fops_tpu_uart_rate);
+	edgetpu_soc_pm_init(etdev);
 
 	if (platform_pwr->after_create)
 		ret = platform_pwr->after_create(etdev);
@@ -649,9 +433,7 @@ static void mobile_pm_before_destroy(struct edgetpu_pm *etpm)
 
 	debugfs_remove_recursive(platform_pwr->debugfs_dir);
 	pm_runtime_disable(etpm->etdev->dev);
-	mobile_pm_cleanup_bts_scenario(etdev);
-	exynos_pm_qos_remove_request(&platform_pwr->int_min);
-	exynos_pm_qos_remove_request(&platform_pwr->mif_min);
+	edgetpu_soc_pm_exit(etdev);
 }
 
 static struct edgetpu_pm_handlers mobile_pm_handlers = {
@@ -661,92 +443,12 @@ static struct edgetpu_pm_handlers mobile_pm_handlers = {
 	.power_down = mobile_power_down,
 };
 
-int mobile_pm_create(struct edgetpu_dev *etdev)
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev)
 {
 	return edgetpu_pm_create(etdev, &mobile_pm_handlers);
 }
 
-void mobile_pm_destroy(struct edgetpu_dev *etdev)
+void edgetpu_mobile_pm_destroy(struct edgetpu_dev *etdev)
 {
 	edgetpu_pm_destroy(etdev);
 }
-
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
-{
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
-	s32 int_val = (pm_qos_val >> PM_QOS_INT_SHIFT) * PM_QOS_FACTOR;
-	s32 mif_val = (pm_qos_val & PM_QOS_MIF_MASK) * PM_QOS_FACTOR;
-
-	etdev_dbg(etdev, "%s: pm_qos request - int = %d mif = %d\n", __func__, int_val, mif_val);
-
-	exynos_pm_qos_update_request(&platform_pwr->int_min, int_val);
-	exynos_pm_qos_update_request(&platform_pwr->mif_min, mif_val);
-}
-
-static void mobile_pm_activate_bts_scenario(struct edgetpu_dev *etdev)
-{
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
-	int performance_scenario = platform_pwr->performance_scenario;
-
-	/* bts_add_scenario() keeps track of reference count internally.*/
-	int ret;
-
-	if (!performance_scenario)
-		return;
-	mutex_lock(&platform_pwr->scenario_lock);
-	ret = bts_add_scenario(performance_scenario);
-	if (ret)
-		etdev_warn_once(etdev, "error %d adding BTS scenario %u\n", ret,
-				performance_scenario);
-	else
-		platform_pwr->scenario_count++;
-
-	etdev_dbg(etdev, "BTS Scenario activated: %d\n", platform_pwr->scenario_count);
-	mutex_unlock(&platform_pwr->scenario_lock);
-}
-
-static void mobile_pm_deactivate_bts_scenario(struct edgetpu_dev *etdev)
-{
-	/* bts_del_scenario() keeps track of reference count internally.*/
-	int ret;
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
-	int performance_scenario = platform_pwr->performance_scenario;
-
-	if (!performance_scenario)
-		return;
-	mutex_lock(&platform_pwr->scenario_lock);
-	if (!platform_pwr->scenario_count) {
-		etdev_warn(etdev, "Unbalanced bts deactivate\n");
-		mutex_unlock(&platform_pwr->scenario_lock);
-		return;
-	}
-	ret = bts_del_scenario(performance_scenario);
-	if (ret)
-		etdev_warn_once(etdev, "error %d deleting BTS scenario %u\n", ret,
-				performance_scenario);
-	else
-		platform_pwr->scenario_count--;
-
-	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n", platform_pwr->scenario_count);
-	mutex_unlock(&platform_pwr->scenario_lock);
-}
-
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
-{
-	etdev_dbg(etdev, "%s: bts request - val = %u\n", __func__, bts_val);
-
-	switch (bts_val) {
-	case 0:
-		mobile_pm_deactivate_bts_scenario(etdev);
-		break;
-	case 1:
-		mobile_pm_activate_bts_scenario(etdev);
-		break;
-	default:
-		etdev_warn(etdev, "%s: invalid BTS request value: %u\n", __func__, bts_val);
-		break;
-	}
-}
diff --git a/drivers/edgetpu/mobile-pm.h b/drivers/edgetpu/mobile-pm.h
index 5f487bc..b9b86b4 100644
--- a/drivers/edgetpu/mobile-pm.h
+++ b/drivers/edgetpu/mobile-pm.h
@@ -9,55 +9,17 @@
 #define __MOBILE_PM_H__
 
 #include "edgetpu-internal.h"
-#include "edgetpu-kci.h"
-
-/*
- * Request codes from firmware
- * Values must match with firmware code base
- */
-enum mobile_reverse_kci_code {
-	RKCI_CODE_PM_QOS = GCIP_RKCI_CHIP_CODE_FIRST + 1,
-	RKCI_CODE_BTS = GCIP_RKCI_CHIP_CODE_FIRST + 2,
-	/* The above codes have been deprecated. */
-
-	RKCI_CODE_PM_QOS_BTS = GCIP_RKCI_CHIP_CODE_FIRST + 3,
-};
-
-#define MAX_VOLTAGE_VAL 1250000
-#define TPU_DEBUG_REQ (1 << 31)
-#define TPU_VDD_TPU_DEBUG (0 << 27)
-#define TPU_VDD_TPU_M_DEBUG (1 << 27)
-#define TPU_VDD_INT_M_DEBUG (2 << 27)
-#define TPU_CLK_CORE_DEBUG (3 << 27)
-#define TPU_CLK_CTL_DEBUG (4 << 27)
-#define TPU_CLK_AXI_DEBUG (5 << 27)
-#define TPU_CLK_APB_DEBUG (6 << 27)
-#define TPU_CLK_UART_DEBUG (7 << 27)
-#define TPU_CORE_PWR_DEBUG (8 << 27)
-#define TPU_DEBUG_VALUE_MASK ((1 << 27) - 1)
 
 /*
  * Initialize a power management interface for an edgetpu device on mobile
  * chipsets.
  * Needs to be called after the devices's platform_pwr struct has been initialized.
  */
-int mobile_pm_create(struct edgetpu_dev *etdev);
-
-/*
- * Wrapper for chip-specific implementation.
- * Typically calls mobile_pm_create after initializing the platform_pwr struct.
- */
-int edgetpu_chip_pm_create(struct edgetpu_dev *etdev);
+int edgetpu_mobile_pm_create(struct edgetpu_dev *etdev);
 
 /*
  * Destroy power management interface for an edgetpu device on mobile chipsets.
  */
-void mobile_pm_destroy(struct edgetpu_dev *etdev);
-
-/* Set required QoS value for the edgetpu device. */
-void mobile_pm_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val);
-
-/* Set BTS value for the edgetpu device. */
-void mobile_pm_set_bts(struct edgetpu_dev *etdev, u16 bts_val);
+void edgetpu_mobile_pm_destroy(struct edgetpu_dev *etdev);
 
 #endif /* __MOBILE_PM_H__ */
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index a969f20..b37f618 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -8,63 +8,123 @@
 #include <linux/acpm_dvfs.h>
 #include <linux/device.h>
 #include <linux/gsa/gsa_tpu.h>
+#include <linux/platform_device.h>
 #include <linux/thermal.h>
 #include <linux/types.h>
+#include <soc/google/bts.h>
+#include <soc/google/exynos_pm_qos.h>
 #include <soc/google/gs_tmu.h>
 
 #include "edgetpu-internal.h"
 #include "edgetpu-firmware.h"
+#include "edgetpu-kci.h"
 #include "edgetpu-mobile-platform.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-thermal.h"
 #include "mobile-firmware.h"
-
+#include "mobile-soc-gsx01.h"
 
 /* TODO(b/199681752): check whether this domain is still correct for Rio */
 #define TPU_ACPM_DOMAIN	7
 
-#define SSMT_NS_READ_STREAM_VID_OFFSET(n) (0x1000u + (0x4u * (n)))
-#define SSMT_NS_WRITE_STREAM_VID_OFFSET(n) (0x1200u + (0x4u * (n)))
+#define MAX_VOLTAGE_VAL 1250000
+
+#define TPU_DEBUG_REQ (1 << 31)
+
+#define TPU_DEBUG_VALUE_SHIFT (27)
+#define TPU_DEBUG_VALUE_MASK ((1 << TPU_DEBUG_VALUE_SHIFT) - 1)
+#define TPU_VDD_TPU_DEBUG (0 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_VDD_TPU_M_DEBUG (1 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_VDD_INT_M_DEBUG (2 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CLK_CORE_DEBUG (3 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CLK_CTL_DEBUG (4 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CLK_AXI_DEBUG (5 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CLK_APB_DEBUG (6 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CLK_UART_DEBUG (7 << TPU_DEBUG_VALUE_SHIFT)
+#define TPU_CORE_PWR_DEBUG (8 << TPU_DEBUG_VALUE_SHIFT)
+
+/*
+ * Encode INT/MIF values as a 16 bit pair in the 32-bit return value
+ * (in units of MHz, to provide enough range)
+ */
+#define PM_QOS_INT_SHIFT                (16)
+#define PM_QOS_MIF_MASK                 (0xFFFF)
+#define PM_QOS_FACTOR                   (1000)
+
+#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))
+
+#define SSMT_BYPASS	(1 << 31)
+
+/* TODO(b/240363978): Remove once ACPM is ready. */
+#if !IS_ENABLED(CONFIG_EDGETPU_TEST)
+unsigned long exynos_acpm_rate = 0;
+#endif /* IS_ENABLED(CONFIG_EDGETPU_TEST) */
+
+static int gsx01_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
+{
+	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
+	struct platform_device *pdev = to_platform_device(etdev->dev);
+	struct resource *res;
+	int ret, i;
+	void __iomem *ssmt_base;
+	char ssmt_name[] = "ssmt_d0";
+
+	etmdev->ssmt_base =
+		devm_kcalloc(etdev->dev, etdev->num_ssmts, sizeof(*etmdev->ssmt_base), GFP_KERNEL);
+
+	if (!etmdev->ssmt_base)
+		return -ENOMEM;
+
+	for (i = 0; i < etdev->num_ssmts; i++) {
+		sprintf(ssmt_name, "ssmt_d%d", i);
+
+		res = platform_get_resource_byname(pdev, IORESOURCE_MEM, ssmt_name);
+		if (!res) {
+			etdev_warn(etdev, "Failed to find SSMT_D%d register base", i);
+			return -EINVAL;
+		}
+		ssmt_base = devm_ioremap_resource(&pdev->dev, res);
+		if (IS_ERR(ssmt_base)) {
+			ret = PTR_ERR(ssmt_base);
+			etdev_warn(etdev, "Failed to map SSMT_D%d register base: %d", i, ret);
+			return ret;
+		}
+		etmdev->ssmt_base[i] = ssmt_base;
+	}
+
+	return 0;
+}
+
+void edgetpu_soc_init(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	int ret;
 
-#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))
+	ret = gsx01_parse_ssmt(etmdev);
+	if (ret)
+		dev_warn(etdev->dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
+}
 
 static void gsx01_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);
+	int i, j;
 
-	/*
-	 * 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;
+	for (i = 0; i < etdev->num_ssmts; i++)
+		if (!etmdev->ssmt_base[i])
+			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));
+	for (i = 0; i < etdev->num_ssmts; i++) {
+		etdev_dbg(etdev, "Setting up SSMT_D%d to feed-through mode\n", i);
+
+		for (j = 0; j < EDGETPU_MAX_STREAM_ID; j++) {
+			writel(SSMT_BYPASS, SSMT_NS_READ_STREAM_VID_REG(etmdev->ssmt_base[i], j));
+			writel(SSMT_BYPASS, SSMT_NS_WRITE_STREAM_VID_REG(etmdev->ssmt_base[i], j));
+		}
 	}
 }
 
@@ -74,6 +134,132 @@ int edgetpu_soc_prepare_firmware(struct edgetpu_dev *etdev)
 	return 0;
 }
 
+static void gsx01_cleanup_bts_scenario(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	if (!performance_scenario)
+		return;
+
+	mutex_lock(&platform_pwr->scenario_lock);
+	while (platform_pwr->scenario_count) {
+		int ret = bts_del_scenario(performance_scenario);
+
+		if (ret) {
+			platform_pwr->scenario_count = 0;
+			etdev_warn_once(etdev, "error %d in cleaning up BTS scenario %u\n", ret,
+					performance_scenario);
+			break;
+		}
+		platform_pwr->scenario_count--;
+	}
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+static void gsx01_activate_bts_scenario(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	/* bts_add_scenario() keeps track of reference count internally.*/
+	int ret;
+
+	if (!performance_scenario)
+		return;
+	mutex_lock(&platform_pwr->scenario_lock);
+	ret = bts_add_scenario(performance_scenario);
+	if (ret)
+		etdev_warn_once(etdev, "error %d adding BTS scenario %u\n", ret,
+				performance_scenario);
+	else
+		platform_pwr->scenario_count++;
+
+	etdev_dbg(etdev, "BTS Scenario activated: %d\n", platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+static void gsx01_deactivate_bts_scenario(struct edgetpu_dev *etdev)
+{
+	/* bts_del_scenario() keeps track of reference count internally.*/
+	int ret;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	int performance_scenario = platform_pwr->performance_scenario;
+
+	if (!performance_scenario)
+		return;
+	mutex_lock(&platform_pwr->scenario_lock);
+	if (!platform_pwr->scenario_count) {
+		etdev_warn(etdev, "Unbalanced bts deactivate\n");
+		mutex_unlock(&platform_pwr->scenario_lock);
+		return;
+	}
+	ret = bts_del_scenario(performance_scenario);
+	if (ret)
+		etdev_warn_once(etdev, "error %d deleting BTS scenario %u\n", ret,
+				performance_scenario);
+	else
+		platform_pwr->scenario_count--;
+
+	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n", platform_pwr->scenario_count);
+	mutex_unlock(&platform_pwr->scenario_lock);
+}
+
+static void gsx01_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
+{
+	etdev_dbg(etdev, "%s: bts request - val = %u\n", __func__, bts_val);
+
+	switch (bts_val) {
+	case 0:
+		gsx01_deactivate_bts_scenario(etdev);
+		break;
+	case 1:
+		gsx01_activate_bts_scenario(etdev);
+		break;
+	default:
+		etdev_warn(etdev, "%s: invalid BTS request value: %u\n", __func__, bts_val);
+		break;
+	}
+}
+
+static void gsx01_set_pm_qos(struct edgetpu_dev *etdev, u32 pm_qos_val)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+	s32 int_val = (pm_qos_val >> PM_QOS_INT_SHIFT) * PM_QOS_FACTOR;
+	s32 mif_val = (pm_qos_val & PM_QOS_MIF_MASK) * PM_QOS_FACTOR;
+
+	etdev_dbg(etdev, "%s: pm_qos request - int = %d mif = %d\n", __func__, int_val, mif_val);
+
+	exynos_pm_qos_update_request(&platform_pwr->int_min, int_val);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, mif_val);
+}
+
+void edgetpu_soc_handle_reverse_kci(struct edgetpu_dev *etdev,
+				    struct gcip_kci_response_element *resp)
+{
+	int ret;
+
+	switch (resp->code) {
+	case RKCI_CODE_PM_QOS_BTS:
+		/* FW indicates to ignore the request by setting them to undefined values. */
+		if (resp->retval != (typeof(resp->retval))~0ull)
+			gsx01_set_pm_qos(etdev, resp->retval);
+		if (resp->status != (typeof(resp->status))~0ull)
+			gsx01_set_bts(etdev, resp->status);
+		ret = edgetpu_kci_resp_rkci_ack(etdev, resp);
+		if (ret)
+			etdev_err(etdev, "failed to send rkci resp for %llu (%d)", resp->seq, ret);
+		break;
+	default:
+		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n", __func__, resp->code);
+		break;
+	}
+}
+
 long edgetpu_soc_pm_get_rate(int flags)
 {
 	return exynos_acpm_get_rate(TPU_ACPM_DOMAIN, flags);
@@ -94,6 +280,195 @@ int edgetpu_soc_pm_set_policy(u64 val)
 	return exynos_acpm_set_policy(TPU_ACPM_DOMAIN, val);
 }
 
+static int edgetpu_core_rate_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG);
+	return 0;
+}
+
+static int edgetpu_core_rate_set(void *data, u64 val)
+{
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CORE_DEBUG;
+	dbg_rate_req |= val;
+
+	return edgetpu_soc_pm_set_rate(dbg_rate_req);
+}
+
+static int edgetpu_ctl_rate_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG);
+	return 0;
+}
+
+static int edgetpu_ctl_rate_set(void *data, u64 val)
+{
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_CTL_DEBUG;
+	dbg_rate_req |= 1000;
+
+	return edgetpu_soc_pm_set_rate(dbg_rate_req);
+}
+
+static int edgetpu_axi_rate_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG);
+	return 0;
+}
+
+static int edgetpu_axi_rate_set(void *data, u64 val)
+{
+	unsigned long dbg_rate_req;
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_CLK_AXI_DEBUG;
+	dbg_rate_req |= 1000;
+
+	return edgetpu_soc_pm_set_rate(dbg_rate_req);
+}
+
+static int edgetpu_apb_rate_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_APB_DEBUG);
+	return 0;
+}
+
+static int edgetpu_uart_rate_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_CLK_UART_DEBUG);
+	return 0;
+}
+
+static int edgetpu_vdd_int_m_set(void *data, u64 val)
+{
+	struct device *dev = (struct device *)data;
+	unsigned long dbg_rate_req;
+
+	if (val > MAX_VOLTAGE_VAL) {
+		dev_err(dev, "Preventing INT_M voltage > %duV", MAX_VOLTAGE_VAL);
+		return -EINVAL;
+	}
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG;
+	dbg_rate_req |= val;
+
+	return edgetpu_soc_pm_set_rate(dbg_rate_req);
+}
+
+static int edgetpu_vdd_int_m_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_INT_M_DEBUG);
+	return 0;
+}
+
+static int edgetpu_vdd_tpu_set(void *data, u64 val)
+{
+	int ret;
+	struct device *dev = (struct device *)data;
+	unsigned long dbg_rate_req;
+
+	if (val > MAX_VOLTAGE_VAL) {
+		dev_err(dev, "Preventing VDD_TPU voltage > %duV", MAX_VOLTAGE_VAL);
+		return -EINVAL;
+	}
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG;
+	dbg_rate_req |= val;
+
+	ret = edgetpu_soc_pm_set_rate(dbg_rate_req);
+	return ret;
+}
+
+static int edgetpu_vdd_tpu_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_TPU_DEBUG);
+	return 0;
+}
+
+static int edgetpu_vdd_tpu_m_set(void *data, u64 val)
+{
+	int ret;
+	struct device *dev = (struct device *)data;
+	unsigned long dbg_rate_req;
+
+	if (val > MAX_VOLTAGE_VAL) {
+		dev_err(dev, "Preventing VDD_TPU voltage > %duV", MAX_VOLTAGE_VAL);
+		return -EINVAL;
+	}
+
+	dbg_rate_req = TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG;
+	dbg_rate_req |= val;
+
+	ret = edgetpu_soc_pm_set_rate(dbg_rate_req);
+	return ret;
+}
+
+static int edgetpu_vdd_tpu_m_get(void *data, u64 *val)
+{
+	*val = edgetpu_soc_pm_get_rate(TPU_DEBUG_REQ | TPU_VDD_TPU_M_DEBUG);
+	return 0;
+}
+
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_core_rate, edgetpu_core_rate_get, edgetpu_core_rate_set,
+			 "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_ctl_rate, edgetpu_ctl_rate_get, edgetpu_ctl_rate_set, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_axi_rate, edgetpu_axi_rate_get, edgetpu_axi_rate_set, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_apb_rate, edgetpu_apb_rate_get, NULL, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_uart_rate, edgetpu_uart_rate_get, NULL, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_int_m, edgetpu_vdd_int_m_get, edgetpu_vdd_int_m_set,
+			 "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu, edgetpu_vdd_tpu_get, edgetpu_vdd_tpu_set, "%llu\n");
+DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu_m, edgetpu_vdd_tpu_m_get, edgetpu_vdd_tpu_m_set,
+			 "%llu\n");
+
+void edgetpu_soc_pm_power_down(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	/* Remove our vote for INT/MIF state (if any) */
+	exynos_pm_qos_update_request(&platform_pwr->int_min, 0);
+	exynos_pm_qos_update_request(&platform_pwr->mif_min, 0);
+
+	gsx01_cleanup_bts_scenario(etdev);
+}
+
+int edgetpu_soc_pm_init(struct edgetpu_dev *etdev)
+{
+	struct device *dev = etdev->dev;
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	exynos_pm_qos_add_request(&platform_pwr->int_min, PM_QOS_DEVICE_THROUGHPUT, 0);
+	exynos_pm_qos_add_request(&platform_pwr->mif_min, PM_QOS_BUS_THROUGHPUT, 0);
+
+	platform_pwr->performance_scenario = bts_get_scenindex("tpu_performance");
+	if (!platform_pwr->performance_scenario)
+		dev_warn(etdev->dev, "tpu_performance BTS scenario not found\n");
+	platform_pwr->scenario_count = 0;
+
+	debugfs_create_file("vdd_tpu", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_tpu);
+	debugfs_create_file("vdd_tpu_m", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_tpu_m);
+	debugfs_create_file("vdd_int_m", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_vdd_int_m);
+	debugfs_create_file("core_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_core_rate);
+	debugfs_create_file("ctl_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_ctl_rate);
+	debugfs_create_file("axi_rate", 0660, platform_pwr->debugfs_dir, dev, &fops_tpu_axi_rate);
+	debugfs_create_file("apb_rate", 0440, platform_pwr->debugfs_dir, dev, &fops_tpu_apb_rate);
+	debugfs_create_file("uart_rate", 0440, platform_pwr->debugfs_dir, dev, &fops_tpu_uart_rate);
+	return 0;
+}
+
+void edgetpu_soc_pm_exit(struct edgetpu_dev *etdev)
+{
+	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
+
+	gsx01_cleanup_bts_scenario(etdev);
+	exynos_pm_qos_remove_request(&platform_pwr->int_min);
+	exynos_pm_qos_remove_request(&platform_pwr->mif_min);
+}
+
 static int tpu_pause_callback(enum thermal_pause_state action, void *dev)
 {
 	int ret = -EINVAL;
diff --git a/drivers/edgetpu/mobile-soc-gsx01.h b/drivers/edgetpu/mobile-soc-gsx01.h
new file mode 100644
index 0000000..5e56dca
--- /dev/null
+++ b/drivers/edgetpu/mobile-soc-gsx01.h
@@ -0,0 +1,23 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * Definitions for GSx01 SoCs.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#ifndef __MOBILE_SOC_GSX01_H__
+#define __MOBILE_SOC_GSX01_H__
+
+/*
+ * Request codes from firmware
+ * Values must match with firmware code base
+ */
+enum gsx01_reverse_kci_code {
+	RKCI_CODE_PM_QOS = GCIP_RKCI_CHIP_CODE_FIRST + 1,
+	RKCI_CODE_BTS = GCIP_RKCI_CHIP_CODE_FIRST + 2,
+	/* The above codes have been deprecated. */
+
+	RKCI_CODE_PM_QOS_BTS = GCIP_RKCI_CHIP_CODE_FIRST + 3,
+};
+
+#endif /* __MOBILE_SOC_GSX01_H__ */
diff --git a/drivers/edgetpu/rio-device.c b/drivers/edgetpu/rio-device.c
index 984e76a..bf9e1be 100644
--- a/drivers/edgetpu/rio-device.c
+++ b/drivers/edgetpu/rio-device.c
@@ -67,28 +67,6 @@ void edgetpu_mark_probe_fail(struct edgetpu_dev *etdev)
 {
 }
 
-void edgetpu_chip_handle_reverse_kci(struct edgetpu_dev *etdev,
-				     struct gcip_kci_response_element *resp)
-{
-	int ret;
-
-	switch (resp->code) {
-	case RKCI_CODE_PM_QOS_BTS:
-		/* FW indicates to ignore the request by setting them to undefined values. */
-		if (resp->retval != (typeof(resp->retval))~0ull)
-			mobile_pm_set_pm_qos(etdev, resp->retval);
-		if (resp->status != (typeof(resp->status))~0ull)
-			mobile_pm_set_bts(etdev, resp->status);
-		ret = edgetpu_kci_resp_rkci_ack(etdev, resp);
-		if (ret)
-			etdev_err(etdev, "failed to send rkci resp for %llu (%d)", resp->seq, ret);
-		break;
-	default:
-		etdev_warn(etdev, "%s: Unrecognized KCI request: %u\n", __func__, resp->code);
-		break;
-	}
-}
-
 int edgetpu_chip_get_ext_mailbox_index(u32 mbox_type, u32 *start, u32 *end)
 {
 	switch (mbox_type) {
diff --git a/drivers/edgetpu/rio-pm.c b/drivers/edgetpu/rio-pm.c
index 28fadb1..aab8fa5 100644
--- a/drivers/edgetpu/rio-pm.c
+++ b/drivers/edgetpu/rio-pm.c
@@ -168,5 +168,5 @@ int edgetpu_chip_pm_create(struct edgetpu_dev *etdev)
 	platform_pwr->lpm_down = rio_lpm_down;
 	platform_pwr->block_down = rio_block_down;
 
-	return mobile_pm_create(etdev);
+	return edgetpu_mobile_pm_create(etdev);
 }
diff --git a/drivers/edgetpu/rio/config.h b/drivers/edgetpu/rio/config.h
index 12eec6f..d358953 100644
--- a/drivers/edgetpu/rio/config.h
+++ b/drivers/edgetpu/rio/config.h
@@ -12,6 +12,10 @@
 
 #define EDGETPU_NUM_CORES 2
 
+#define EDGETPU_NUM_SSMTS 2
+
+#define EDGETPU_MAX_STREAM_ID 64
+
 /* 1 context per VII/group plus 1 for KCI */
 #define EDGETPU_NCONTEXTS 16
 /* Max number of virtual context IDs that can be allocated for one device. */
@@ -22,9 +26,6 @@
 /* Placeholder value */
 #define EDGETPU_TZ_MAILBOX_ID 31
 
-/* Responds to PMQoS-BTS RKCI */
-#define EDGETPU_FEATURE_RKCI_RESPONSE
-
 /*
  * Size of the area in remapped DRAM reserved for firmware code and internal
  * data. This must match the firmware's linker file.