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

edgetpu: Fix bcl_dev referencing
edgetpu: Remove unnecessary telemetry check
gcip: use pad/pap instead of llx for logging
edgetpu: Remove the non caller mem path in telemetry
Bug: 239374826
edgetpu: fix invalid paging on invalid map size
Bug: 245095356
edgetpu: Set TPU_ACPM_DOMAIN to 9 for Rio
Bug: 199681752
edgetpu: move soc-specific data to new struct
edgetpu: move edgetpu_soc_init call to edgetpu core

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

13 files changed, 110 insertions, 192 deletions

diff --git a/drivers/edgetpu/edgetpu-config.h b/drivers/edgetpu/edgetpu-config.h
index c496a33..9aa87bb 100644
--- a/drivers/edgetpu/edgetpu-config.h
+++ b/drivers/edgetpu/edgetpu-config.h
@@ -25,10 +25,6 @@
 #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 dc8ba6c..43ef85a 100644
--- a/drivers/edgetpu/edgetpu-core.c
+++ b/drivers/edgetpu/edgetpu-core.c
@@ -29,6 +29,7 @@
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 #include "edgetpu-mmu.h"
+#include "edgetpu-soc.h"
 #include "edgetpu-sw-watchdog.h"
 #include "edgetpu-telemetry.h"
 #include "edgetpu-usage-stats.h"
@@ -430,6 +431,9 @@ int edgetpu_device_add(struct edgetpu_dev *etdev,
 	etdev->state = ETDEV_STATE_NOFW;
 	etdev->freq_count = 0;
 	mutex_init(&etdev->freq_lock);
+	ret = edgetpu_soc_init(etdev);
+	if (ret)
+		return ret;
 
 	ret = edgetpu_fs_add(etdev, iface_params, num_ifaces);
 	if (ret) {
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index fc0ff01..91e83d7 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -655,6 +655,8 @@ static struct page **edgetpu_pin_user_pages(struct edgetpu_device_group *group,
 	}
 	offset = host_addr & (PAGE_SIZE - 1);
 	num_pages = DIV_ROUND_UP((size + offset), PAGE_SIZE);
+	if (num_pages * PAGE_SIZE < size + offset)
+		return ERR_PTR(-EINVAL);
 	etdev_dbg(etdev, "%s: hostaddr=%#llx pages=%u", __func__, host_addr, num_pages);
 	/*
 	 * "num_pages" is decided from user-space arguments, don't show warnings
diff --git a/drivers/edgetpu/edgetpu-internal.h b/drivers/edgetpu/edgetpu-internal.h
index a9260e7..5dc4723 100644
--- a/drivers/edgetpu/edgetpu-internal.h
+++ b/drivers/edgetpu/edgetpu-internal.h
@@ -87,6 +87,7 @@ struct edgetpu_coherent_mem {
 struct edgetpu_device_group;
 struct edgetpu_wakelock;
 struct edgetpu_dev_iface;
+struct edgetpu_soc_data;
 
 #define EDGETPU_NUM_PERDIE_EVENTS	2
 #define perdie_event_id_to_num(event_id)				      \
@@ -170,7 +171,6 @@ 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
@@ -178,6 +178,8 @@ struct edgetpu_dev {
 	struct edgetpu_dev_iface *etiface;
 	char dev_name[EDGETPU_DEVICE_NAME_MAX];
 	struct edgetpu_mapped_resource regs; /* ioremapped CSRs */
+	/* SoC-specific data */
+	struct edgetpu_soc_data *soc_data;
 	struct dentry *d_entry;    /* debugfs dir for this device */
 	struct mutex state_lock;   /* protects state of this device */
 	enum edgetpu_dev_state state;
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.c b/drivers/edgetpu/edgetpu-mobile-platform.c
index ce7e75c..f74e9b1 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.c
+++ b/drivers/edgetpu/edgetpu-mobile-platform.c
@@ -309,7 +309,6 @@ 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)) {
@@ -325,8 +324,6 @@ 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;
 
diff --git a/drivers/edgetpu/edgetpu-mobile-platform.h b/drivers/edgetpu/edgetpu-mobile-platform.h
index afdaa9d..02d560f 100644
--- a/drivers/edgetpu/edgetpu-mobile-platform.h
+++ b/drivers/edgetpu/edgetpu-mobile-platform.h
@@ -14,11 +14,6 @@
 #include <linux/kernel.h>
 #include <linux/mutex.h>
 #include <linux/types.h>
-#include <soc/google/exynos_pm_qos.h>
-
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-#include <soc/google/bcl.h>
-#endif
 
 #include "edgetpu-config.h"
 #include "edgetpu-internal.h"
@@ -33,13 +28,6 @@ struct edgetpu_mobile_platform_pwr {
 	struct mutex state_lock;
 	u64 min_state;
 	u64 requested_state;
-	/* INT/MIF requests for memory bandwidth */
-	struct exynos_pm_qos_request int_min;
-	struct exynos_pm_qos_request mif_min;
-	/* BTS */
-	unsigned int performance_scenario;
-	int scenario_count;
-	struct mutex scenario_lock;
 
 	/* LPM callbacks, NULL for chips without LPM */
 	int (*lpm_up)(struct edgetpu_dev *etdev);
@@ -79,15 +67,10 @@ struct edgetpu_mobile_platform_dev {
 	 * May be NULL if the chip does not support firmware authentication
 	 */
 	struct device *gsa_dev;
-	/* Virtual address of the SSMT block for this chip. */
-	void __iomem **ssmt_base;
 	/* Coherent log buffer */
 	struct edgetpu_coherent_mem *log_mem;
 	/* Coherent trace buffer */
 	struct edgetpu_coherent_mem *trace_mem;
-#if IS_ENABLED(CONFIG_GOOGLE_BCL)
-	struct bcl_device *bcl_dev;
-#endif
 	/* subsystem coredump info struct */
 	struct mobile_sscd_info sscd_info;
 	/* Protects TZ Mailbox client pointer */
diff --git a/drivers/edgetpu/edgetpu-soc.h b/drivers/edgetpu/edgetpu-soc.h
index 5e87676..00e91ef 100644
--- a/drivers/edgetpu/edgetpu-soc.h
+++ b/drivers/edgetpu/edgetpu-soc.h
@@ -17,7 +17,7 @@
 /* SoC-specific calls for the following functions. */
 
 /* Probe-time init */
-void edgetpu_soc_init(struct edgetpu_dev *etdev);
+int edgetpu_soc_init(struct edgetpu_dev *etdev);
 
 /* Prep for running firmware: set access control, etc. */
 int edgetpu_soc_prepare_firmware(struct edgetpu_dev *etdev);
diff --git a/drivers/edgetpu/edgetpu-telemetry.c b/drivers/edgetpu/edgetpu-telemetry.c
index ca90e72..ce4739f 100644
--- a/drivers/edgetpu/edgetpu-telemetry.c
+++ b/drivers/edgetpu/edgetpu-telemetry.c
@@ -45,8 +45,6 @@ static int telemetry_kci(struct edgetpu_dev *etdev, struct edgetpu_telemetry *te
 {
 	int err;
 
-	if (!tel->inited)
-		return -ENODEV;
 	etdev_dbg(etdev, "Sending KCI %s", tel->name);
 	err = send_kci(etdev->etkci, tel->coherent_mem.tpu_addr, tel->coherent_mem.size);
 
@@ -69,8 +67,6 @@ static int telemetry_set_event(struct edgetpu_dev *etdev,
 	struct eventfd_ctx *ctx;
 	ulong flags;
 
-	if (!tel->inited)
-		return -ENODEV;
 	ctx = eventfd_ctx_fdget(eventfd);
 	if (IS_ERR(ctx))
 		return PTR_ERR(ctx);
@@ -89,8 +85,6 @@ static void telemetry_unset_event(struct edgetpu_dev *etdev,
 {
 	ulong flags;
 
-	if (!tel->inited)
-		return;
 	write_lock_irqsave(&tel->ctx_lock, flags);
 	if (tel->ctx)
 		eventfd_ctx_put(tel->ctx);
@@ -223,8 +217,6 @@ static void telemetry_worker(struct work_struct *work)
 static void telemetry_irq_handler(struct edgetpu_dev *etdev,
 				  struct edgetpu_telemetry *tel)
 {
-	if (!tel->inited)
-		return;
 	spin_lock(&tel->state_lock);
 
 	if (tel->state == EDGETPU_TELEMETRY_ENABLED &&
@@ -238,9 +230,6 @@ static void telemetry_irq_handler(struct edgetpu_dev *etdev,
 static void telemetry_mappings_show(struct edgetpu_telemetry *tel,
 				    struct seq_file *s)
 {
-	if (!tel->inited)
-		return;
-
 	seq_printf(s, "  %#llx %lu %s %#llx %pad\n",
 		   tel->coherent_mem.tpu_addr,
 		   DIV_ROUND_UP(tel->coherent_mem.size, PAGE_SIZE), tel->name,
@@ -249,9 +238,6 @@ static void telemetry_mappings_show(struct edgetpu_telemetry *tel,
 
 static void telemetry_inc_mmap_count(struct edgetpu_telemetry *tel, int dif)
 {
-	if (!tel->inited)
-		return;
-
 	mutex_lock(&tel->mmap_lock);
 	tel->mmapped_count += dif;
 	mutex_unlock(&tel->mmap_lock);
@@ -263,9 +249,6 @@ static int telemetry_mmap_buffer(struct edgetpu_dev *etdev,
 {
 	int ret;
 
-	if (!tel->inited)
-		return -ENODEV;
-
 	mutex_lock(&tel->mmap_lock);
 
 	if (!tel->mmapped_count) {
@@ -290,42 +273,12 @@ static int telemetry_init(struct edgetpu_dev *etdev, struct edgetpu_telemetry *t
 			  const char *name, struct edgetpu_coherent_mem *mem, const size_t size,
 			  void (*fallback)(struct edgetpu_telemetry *))
 {
-	const u32 flags = EDGETPU_MMU_32 | EDGETPU_MMU_HOST;
-	void *vaddr;
-	dma_addr_t dma_addr;
-	tpu_addr_t tpu_addr;
-
-	if (mem) {
-		tel->coherent_mem = *mem;
-		vaddr = mem->vaddr;
-		tel->caller_mem = true;
-	} else {
-		vaddr = dmam_alloc_coherent(etdev->dev, size, &dma_addr,
-					    GFP_KERNEL);
-		if (!vaddr)
-			return -ENOMEM;
-		tpu_addr = edgetpu_mmu_tpu_map(etdev, dma_addr, size,
-					       DMA_BIDIRECTIONAL,
-					       EDGETPU_CONTEXT_KCI, flags);
-		if (!tpu_addr) {
-			dev_err(etdev->dev,
-				"%s: failed to map buffer for '%s'\n",
-				etdev->dev_name, name);
-			return -ENOSPC;
-		}
-		tel->coherent_mem.vaddr = vaddr;
-		tel->coherent_mem.dma_addr = dma_addr;
-		tel->coherent_mem.tpu_addr = tpu_addr;
-		tel->coherent_mem.size = size;
-		tel->caller_mem = false;
-		edgetpu_x86_coherent_mem_set_uc(&tel->coherent_mem);
-	}
-
 	rwlock_init(&tel->ctx_lock);
 	tel->name = name;
 	tel->etdev = etdev;
+	tel->coherent_mem = *mem;
 
-	tel->header = (struct edgetpu_telemetry_header *)vaddr;
+	tel->header = (struct edgetpu_telemetry_header *)mem->vaddr;
 	tel->header->head = 0;
 	tel->header->size = 0;
 	tel->header->tail = 0;
@@ -337,7 +290,6 @@ static int telemetry_init(struct edgetpu_dev *etdev, struct edgetpu_telemetry *t
 	INIT_WORK(&tel->work, telemetry_worker);
 	tel->fallback_fn = fallback;
 	tel->state = EDGETPU_TELEMETRY_ENABLED;
-	tel->inited = true;
 	mutex_init(&tel->mmap_lock);
 	tel->mmapped_count = 0;
 
@@ -349,21 +301,12 @@ static void telemetry_exit(struct edgetpu_dev *etdev,
 {
 	ulong flags;
 
-	if (!tel->inited)
-		return;
 	spin_lock_irqsave(&tel->state_lock, flags);
 	/* Prevent racing with the IRQ handler or worker */
 	tel->state = EDGETPU_TELEMETRY_INVALID;
 	spin_unlock_irqrestore(&tel->state_lock, flags);
 	cancel_work_sync(&tel->work);
 
-	if (tel->coherent_mem.tpu_addr && !tel->caller_mem) {
-		edgetpu_mmu_tpu_unmap(etdev, tel->coherent_mem.tpu_addr,
-				      tel->coherent_mem.size,
-				      EDGETPU_CONTEXT_KCI);
-		tel->coherent_mem.tpu_addr = 0;
-		edgetpu_x86_coherent_mem_set_wb(&tel->coherent_mem);
-	}
 	if (tel->ctx)
 		eventfd_ctx_put(tel->ctx);
 	tel->ctx = NULL;
@@ -375,9 +318,6 @@ int edgetpu_telemetry_init(struct edgetpu_dev *etdev,
 {
 	int ret, i;
 
-	if (!etdev->telemetry)
-		return -ENODEV;
-
 	for (i = 0; i < etdev->num_cores; i++) {
 		ret = telemetry_init(etdev, &etdev->telemetry[i].log, "telemetry_log",
 				     log_mem ? &log_mem[i] : NULL,
@@ -403,9 +343,6 @@ void edgetpu_telemetry_exit(struct edgetpu_dev *etdev)
 {
 	int i;
 
-	if (!etdev->telemetry)
-		return;
-
 	for (i = 0; i < etdev->num_cores; i++) {
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
 		telemetry_exit(etdev, &etdev->telemetry[i].trace);
@@ -418,9 +355,6 @@ int edgetpu_telemetry_kci(struct edgetpu_dev *etdev)
 {
 	int ret;
 
-	if (!etdev->telemetry)
-		return -ENODEV;
-
 	/* Core 0 will notify other cores. */
 	ret = telemetry_kci(etdev, &etdev->telemetry[0].log, edgetpu_kci_map_log_buffer);
 	if (ret)
@@ -440,9 +374,6 @@ int edgetpu_telemetry_set_event(struct edgetpu_dev *etdev,
 {
 	int i, ret;
 
-	if (!etdev->telemetry)
-		return -ENODEV;
-
 	for (i = 0; i < etdev->num_cores; i++) {
 		ret = telemetry_set_event(etdev, select_telemetry(&etdev->telemetry[i], type),
 					  eventfd);
@@ -460,9 +391,6 @@ void edgetpu_telemetry_unset_event(struct edgetpu_dev *etdev,
 {
 	int i;
 
-	if (!etdev->telemetry)
-		return;
-
 	for (i = 0; i < etdev->num_cores; i++)
 		telemetry_unset_event(etdev, select_telemetry(&etdev->telemetry[i], type));
 }
@@ -471,9 +399,6 @@ void edgetpu_telemetry_irq_handler(struct edgetpu_dev *etdev)
 {
 	int i;
 
-	if (!etdev->telemetry)
-		return;
-
 	for (i = 0; i < etdev->num_cores; i++) {
 		telemetry_irq_handler(etdev, &etdev->telemetry[i].log);
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
@@ -487,9 +412,6 @@ void edgetpu_telemetry_mappings_show(struct edgetpu_dev *etdev,
 {
 	int i;
 
-	if (!etdev->telemetry)
-		return;
-
 	for (i = 0; i < etdev->num_cores; i++) {
 		telemetry_mappings_show(&etdev->telemetry[i].log, s);
 #if IS_ENABLED(CONFIG_EDGETPU_TELEMETRY_TRACE)
@@ -501,8 +423,6 @@ void edgetpu_telemetry_mappings_show(struct edgetpu_dev *etdev,
 int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
 				  struct vm_area_struct *vma, int core_id)
 {
-	if (!etdev->telemetry)
-		return -ENODEV;
 	return telemetry_mmap_buffer(etdev, select_telemetry(&etdev->telemetry[core_id], type),
 				     vma);
 }
@@ -510,15 +430,11 @@ int edgetpu_mmap_telemetry_buffer(struct edgetpu_dev *etdev, enum edgetpu_teleme
 void edgetpu_telemetry_inc_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
 				      int core_id)
 {
-	if (!etdev->telemetry)
-		return;
 	telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), 1);
 }
 
 void edgetpu_telemetry_dec_mmap_count(struct edgetpu_dev *etdev, enum edgetpu_telemetry_type type,
 				      int core_id)
 {
-	if (!etdev->telemetry)
-		return;
 	telemetry_inc_mmap_count(select_telemetry(&etdev->telemetry[core_id], type), -1);
 }
diff --git a/drivers/edgetpu/edgetpu-telemetry.h b/drivers/edgetpu/edgetpu-telemetry.h
index f849c8c..b703299 100644
--- a/drivers/edgetpu/edgetpu-telemetry.h
+++ b/drivers/edgetpu/edgetpu-telemetry.h
@@ -71,17 +71,10 @@ struct edgetpu_telemetry {
 
 	struct edgetpu_coherent_mem coherent_mem;
 	struct edgetpu_telemetry_header *header;
-	/*
-	 * If coherent_mem buffer is provided by the caller in
-	 * edgetpu_telemetry_init, the caller is responsible for
-	 * releasing/unmapping it.
-	 */
-	bool caller_mem;
 
 	struct eventfd_ctx *ctx; /* signal this to notify the runtime */
 	rwlock_t ctx_lock; /* protects ctx */
 	const char *name; /* for debugging */
-	bool inited; /* whether telemetry_init() succeeded */
 
 	/* Worker for handling data. */
 	struct work_struct work;
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
index af082b6..47c7de0 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-image-config.c
@@ -38,8 +38,8 @@ static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 		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);
+		dev_dbg(parser->dev, "Image config adding IOMMU mapping: %pad -> %pap", &daddr,
+			&paddr);
 
 		if (unlikely(daddr + size <= daddr || paddr + size <= paddr)) {
 			ret = -EOVERFLOW;
@@ -49,8 +49,8 @@ static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 				       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);
+				"Unable to Map: %d dma_addr: %pad phys_addr: %pap size: %#lx\n",
+				ret, &daddr, &paddr, size);
 			goto err;
 		}
 	}
@@ -76,7 +76,7 @@ static void clear_iommu_mappings(struct gcip_image_config_parser *parser,
 	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,
+		dev_dbg(parser->dev, "Image config removing IOMMU mapping: %pad size=%#lx", &daddr,
 			size);
 		parser->ops->unmap(parser->data, daddr, size, GCIP_IMAGE_CONFIG_FLAGS_SECURE);
 	}
@@ -98,8 +98,8 @@ static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 			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);
+		dev_dbg(parser->dev, "Image config adding NS IOMMU mapping: %pad -> %pap", &daddr,
+			&paddr);
 		if (unlikely(daddr + size <= daddr || paddr + size <= paddr)) {
 			ret = -EOVERFLOW;
 			goto err;
@@ -131,8 +131,8 @@ static void clear_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 	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);
+		dev_dbg(parser->dev, "Image config removing NS IOMMU mapping: %pad size=%#lx",
+			&daddr, size);
 		parser->ops->unmap(parser->data, daddr, size, 0);
 	}
 }
diff --git a/drivers/edgetpu/mobile-pm.c b/drivers/edgetpu/mobile-pm.c
index d145c0c..4d19b2a 100644
--- a/drivers/edgetpu/mobile-pm.c
+++ b/drivers/edgetpu/mobile-pm.c
@@ -21,6 +21,7 @@
 #include "edgetpu-pm.h"
 #include "mobile-firmware.h"
 #include "mobile-pm.h"
+#include "mobile-soc-gsx01.h"
 
 #include "edgetpu-pm.c"
 #include "edgetpu-soc.h"
@@ -301,10 +302,10 @@ static int mobile_power_up(struct edgetpu_pm *etpm)
 		mobile_power_down(etpm);
 	} else {
 #if IS_ENABLED(CONFIG_GOOGLE_BCL)
-		if (!etmdev->bcl_dev)
-			etmdev->bcl_dev = google_retrieve_bcl_handle();
-		if (etmdev->bcl_dev)
-			google_init_tpu_ratio(etmdev->bcl_dev);
+		if (!etdev->soc_data->bcl_dev)
+			etdev->soc_data->bcl_dev = google_retrieve_bcl_handle();
+		if (etdev->soc_data->bcl_dev)
+			google_init_tpu_ratio(etdev->soc_data->bcl_dev);
 #endif
 	}
 
@@ -399,7 +400,6 @@ static int mobile_pm_after_create(struct edgetpu_pm *etpm)
 
 	mutex_init(&platform_pwr->policy_lock);
 	mutex_init(&platform_pwr->state_lock);
-	mutex_init(&platform_pwr->scenario_lock);
 
 	ret = mobile_pwr_state_set(etdev, mobile_get_initial_pwr_state(dev));
 	if (ret)
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index 3bcef1d..d69cd32 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -24,8 +24,7 @@
 #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 TPU_ACPM_DOMAIN	9
 
 #define MAX_VOLTAGE_VAL 1250000
 
@@ -68,21 +67,23 @@ 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 edgetpu_soc_data *soc_data = etdev->soc_data;
 	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);
+	soc_data->num_ssmts = EDGETPU_NUM_SSMTS;
+	soc_data->ssmt_base = devm_kcalloc(etdev->dev, soc_data->num_ssmts,
+					   sizeof(*soc_data->ssmt_base), GFP_KERNEL);
 
-	if (!etmdev->ssmt_base)
+	if (!soc_data->ssmt_base)
 		return -ENOMEM;
 
-	if (unlikely(etdev->num_ssmts > 9))
+	if (unlikely(soc_data->num_ssmts > 9))
 		return -EINVAL;
 
-	for (i = 0; i < etdev->num_ssmts; i++) {
+	for (i = 0; i < soc_data->num_ssmts; i++) {
 		sprintf(ssmt_name, "ssmt_d%d", i);
 
 		res = platform_get_resource_byname(pdev, IORESOURCE_MEM, ssmt_name);
@@ -96,37 +97,45 @@ static int gsx01_parse_ssmt(struct edgetpu_mobile_platform_dev *etmdev)
 			etdev_warn(etdev, "Failed to map SSMT_D%d register base: %d", i, ret);
 			return ret;
 		}
-		etmdev->ssmt_base[i] = ssmt_base;
+		soc_data->ssmt_base[i] = ssmt_base;
 	}
-
 	return 0;
 }
 
-void edgetpu_soc_init(struct edgetpu_dev *etdev)
+int edgetpu_soc_init(struct edgetpu_dev *etdev)
 {
+	struct platform_device *pdev = to_platform_device(etdev->dev);
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	int ret;
 
+	etdev->soc_data = devm_kzalloc(&pdev->dev, sizeof(*etdev->soc_data), GFP_KERNEL);
+	if (!etdev->soc_data)
+		return -ENOMEM;
+
+	mutex_init(&etdev->soc_data->scenario_lock);
 	ret = gsx01_parse_ssmt(etmdev);
 	if (ret)
 		dev_warn(etdev->dev, "SSMT setup failed (%d). Context isolation not enforced", ret);
+	return 0;
 }
 
 static void gsx01_setup_ssmt(struct edgetpu_dev *etdev)
 {
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
+	struct edgetpu_soc_data *soc_data = etdev->soc_data;
 	int i, j;
 
-	for (i = 0; i < etdev->num_ssmts; i++)
-		if (!etmdev->ssmt_base[i])
+	for (i = 0; i < soc_data->num_ssmts; i++)
+		if (!soc_data->ssmt_base[i])
 			return;
 
-	for (i = 0; i < etdev->num_ssmts; i++) {
+	for (i = 0; i < soc_data->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));
+			writel(SSMT_BYPASS,
+			       SSMT_NS_READ_STREAM_VID_REG(soc_data->ssmt_base[i], j));
+			writel(SSMT_BYPASS,
+			       SSMT_NS_WRITE_STREAM_VID_REG(soc_data->ssmt_base[i], j));
 		}
 	}
 }
@@ -139,65 +148,62 @@ int edgetpu_soc_prepare_firmware(struct edgetpu_dev *etdev)
 
 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;
+	struct edgetpu_soc_data *soc_data = etdev->soc_data;
+	int performance_scenario = soc_data->performance_scenario;
 
 	if (!performance_scenario)
 		return;
 
-	mutex_lock(&platform_pwr->scenario_lock);
-	while (platform_pwr->scenario_count) {
+	mutex_lock(&soc_data->scenario_lock);
+	while (soc_data->scenario_count) {
 		int ret = bts_del_scenario(performance_scenario);
 
 		if (ret) {
-			platform_pwr->scenario_count = 0;
+			soc_data->scenario_count = 0;
 			etdev_warn_once(etdev, "error %d in cleaning up BTS scenario %u\n", ret,
 					performance_scenario);
 			break;
 		}
-		platform_pwr->scenario_count--;
+		soc_data->scenario_count--;
 	}
-	mutex_unlock(&platform_pwr->scenario_lock);
+	mutex_unlock(&soc_data->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;
+	struct edgetpu_soc_data *soc_data = etdev->soc_data;
+	int performance_scenario = soc_data->performance_scenario;
 
 	/* bts_add_scenario() keeps track of reference count internally.*/
 	int ret;
 
 	if (!performance_scenario)
 		return;
-	mutex_lock(&platform_pwr->scenario_lock);
+	mutex_lock(&soc_data->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++;
+		soc_data->scenario_count++;
 
-	etdev_dbg(etdev, "BTS Scenario activated: %d\n", platform_pwr->scenario_count);
-	mutex_unlock(&platform_pwr->scenario_lock);
+	etdev_dbg(etdev, "BTS Scenario activated: %d\n", soc_data->scenario_count);
+	mutex_unlock(&soc_data->scenario_lock);
 }
 
 static void gsx01_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 {
 	/* bts_del_scenario() keeps track of reference count internally.*/
+	struct edgetpu_soc_data *soc_data = etdev->soc_data;
+	int performance_scenario = soc_data->performance_scenario;
 	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) {
+	mutex_lock(&soc_data->scenario_lock);
+	if (!soc_data->scenario_count) {
 		etdev_warn(etdev, "Unbalanced bts deactivate\n");
-		mutex_unlock(&platform_pwr->scenario_lock);
+		mutex_unlock(&soc_data->scenario_lock);
 		return;
 	}
 	ret = bts_del_scenario(performance_scenario);
@@ -205,10 +211,10 @@ static void gsx01_deactivate_bts_scenario(struct edgetpu_dev *etdev)
 		etdev_warn_once(etdev, "error %d deleting BTS scenario %u\n", ret,
 				performance_scenario);
 	else
-		platform_pwr->scenario_count--;
+		soc_data->scenario_count--;
 
-	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n", platform_pwr->scenario_count);
-	mutex_unlock(&platform_pwr->scenario_lock);
+	etdev_dbg(etdev, "BTS Scenario deactivated: %d\n", soc_data->scenario_count);
+	mutex_unlock(&soc_data->scenario_lock);
 }
 
 static void gsx01_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
@@ -230,15 +236,13 @@ static void gsx01_set_bts(struct edgetpu_dev *etdev, u16 bts_val)
 
 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);
+	exynos_pm_qos_update_request(&etdev->soc_data->int_min, int_val);
+	exynos_pm_qos_update_request(&etdev->soc_data->mif_min, mif_val);
 }
 
 void edgetpu_soc_handle_reverse_kci(struct edgetpu_dev *etdev,
@@ -427,12 +431,9 @@ DEFINE_DEBUGFS_ATTRIBUTE(fops_tpu_vdd_tpu_m, edgetpu_vdd_tpu_m_get, edgetpu_vdd_
 
 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);
+	exynos_pm_qos_update_request(&etdev->soc_data->int_min, 0);
+	exynos_pm_qos_update_request(&etdev->soc_data->mif_min, 0);
 
 	gsx01_cleanup_bts_scenario(etdev);
 }
@@ -443,13 +444,13 @@ int edgetpu_soc_pm_init(struct edgetpu_dev *etdev)
 	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);
+	exynos_pm_qos_add_request(&etdev->soc_data->int_min, PM_QOS_DEVICE_THROUGHPUT, 0);
+	exynos_pm_qos_add_request(&etdev->soc_data->mif_min, PM_QOS_BUS_THROUGHPUT, 0);
 
-	platform_pwr->performance_scenario = bts_get_scenindex("tpu_performance");
-	if (!platform_pwr->performance_scenario)
+	etdev->soc_data->performance_scenario = bts_get_scenindex("tpu_performance");
+	if (!etdev->soc_data->performance_scenario)
 		dev_warn(etdev->dev, "tpu_performance BTS scenario not found\n");
-	platform_pwr->scenario_count = 0;
+	etdev->soc_data->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);
@@ -464,12 +465,9 @@ int edgetpu_soc_pm_init(struct edgetpu_dev *etdev)
 
 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);
+	exynos_pm_qos_remove_request(&etdev->soc_data->int_min);
+	exynos_pm_qos_remove_request(&etdev->soc_data->mif_min);
 }
 
 static int tpu_pause_callback(enum thermal_pause_state action, void *dev)
diff --git a/drivers/edgetpu/mobile-soc-gsx01.h b/drivers/edgetpu/mobile-soc-gsx01.h
index 5e56dca..c402d8d 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.h
+++ b/drivers/edgetpu/mobile-soc-gsx01.h
@@ -8,6 +8,33 @@
 #ifndef __MOBILE_SOC_GSX01_H__
 #define __MOBILE_SOC_GSX01_H__
 
+#include <linux/kernel.h>
+#include <linux/mutex.h>
+#include <linux/types.h>
+#include <soc/google/exynos_pm_qos.h>
+
+#if IS_ENABLED(CONFIG_GOOGLE_BCL)
+#include <soc/google/bcl.h>
+#endif
+
+/* SoC data for GSx01 platforms */
+struct edgetpu_soc_data {
+	/* Virtual address of the SSMT block for this chip. */
+	void __iomem **ssmt_base;
+	/* Number of SSMTs */
+	uint num_ssmts;
+	/* INT/MIF requests for memory bandwidth */
+	struct exynos_pm_qos_request int_min;
+	struct exynos_pm_qos_request mif_min;
+	/* BTS */
+	unsigned int performance_scenario;
+	int scenario_count;
+	struct mutex scenario_lock;
+#if IS_ENABLED(CONFIG_GOOGLE_BCL)
+	struct bcl_device *bcl_dev;
+#endif
+};
+
 /*
  * Request codes from firmware
  * Values must match with firmware code base