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

edgetpu: rio: Add support for kleaf
Bug: 257343887
edgetpu: Unify the build files
edgetpu: Update the trace header file
edgetpu: signal any unsignaled dma fences at device group disband
Bug: 258868303
edgetpu: dma fence debugfs show owning group
Bug: 258868303 (repeat)
edgetpu: dma fence creation optionally associated with device group
Bug: 258868303 (repeat)
edgetpu: ensure TPU is powered before reading CMU
Bug: 261231679
edgetpu: introduce START_OFFLOAD external cmd
Bug: 247923533
edgetpu: adopt tpu_file to external commands
Bug: 247923533 (repeat)

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

13 files changed, 242 insertions, 152 deletions

diff --git a/drivers/edgetpu/BUILD.bazel b/drivers/edgetpu/BUILD.bazel
new file mode 100644
index 0000000..474ccff
--- /dev/null
+++ b/drivers/edgetpu/BUILD.bazel
@@ -0,0 +1,34 @@
+# NOTE: THIS FILE IS EXPERIMENTAL FOR THE BAZEL MIGRATION AND NOT USED FOR
+# YOUR BUILDS CURRENTLY.
+#
+# It is not yet the source of truth for your build. If you're looking to modify
+# the build file, modify the Android.bp file instead. Do *not* modify this file
+# unless you have coordinated with the team managing the Soong to Bazel
+# migration.
+
+load("//build/kernel/kleaf:kernel.bzl", "kernel_module")
+
+kernel_module(
+    name = "edgetpu.rio",
+    srcs = glob([
+        "**/*.c",
+        "**/*.h",
+        "Kbuild",
+        "Makefille",
+        "gcip-kernel-driver/drivers/gcip/Makefile",
+    ]) + [
+        "//private/google-modules/soc/gs:gs_soc_headers",
+    ],
+    outs = [
+        "rio.ko",
+    ],
+    kernel_build = "//private/google-modules/soc/gs:gs_kernel_build",
+    visibility = [
+        "//private/devices/google:__subpackages__",
+        "//private/google-modules/gxp/zuma:__pkg__",
+        "//private/google-modules/soc/gs:__pkg__",
+    ],
+    deps = [
+        "//private/google-modules/soc/gs:gs_soc_module",
+    ],
+)
diff --git a/drivers/edgetpu/Kbuild b/drivers/edgetpu/Kbuild
index 6bd52cf..c9fe855 100644
--- a/drivers/edgetpu/Kbuild
+++ b/drivers/edgetpu/Kbuild
@@ -1,21 +1,37 @@
-obj-m	+= rio.o
+# SPDX-License-Identifier: GPL-2.0
+
 # Use the absolute path of this Makefile to get the source directory.
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
-ccflags-y	+= $(EXTRA_CFLAGS) -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -I$(CURRENT_DIR)/include \
-		   -DCONFIG_GCIP=1 -I$(CURRENT_DIR)/gcip-kernel-driver/include
 GIT_PATH=$(CURRENT_DIR)/../../
-ifeq ($(shell git --git-dir=$(GIT_PATH)/.git rev-parse --is-inside-work-tree),true)
-	GIT_REPO_STATE=$(shell (git --git-dir=$(GIT_PATH)/.git --work-tree=$(GIT_PATH) status --porcelain | grep -q .) && echo -dirty)
-	ccflags-y	+= -DGIT_REPO_TAG=\"$(shell git --git-dir=$(GIT_PATH)/.git rev-parse --short HEAD)$(GIT_REPO_STATE)\"
+ccflags-y	+= -DCONFIG_EDGETPU_TELEMETRY_TRACE=1 -DCONFIG_GCIP=1 \
+		   -I$(CURRENT_DIR)/include \
+		   -I$(CURRENT_DIR)/gcip-kernel-driver/include
+
+ifeq ($(shell git -C $(GIT_PATH) rev-parse --is-inside-work-tree),true)
+	GIT_REPO_STATE=$(shell (git -C $(GIT_PATH) status --porcelain | grep -q .) && echo -dirty)
+	ccflags-y	+= -DGIT_REPO_TAG=\"$(shell git -C $(GIT_PATH) rev-parse --short HEAD)$(GIT_REPO_STATE)\"
 else
 	ccflags-y	+= -DGIT_REPO_TAG=\"Not\ a\ git\ repository\"
 endif
 
+obj-$(CONFIG_RIO)	+= rio.o
+
 GCIP_DIR=gcip-kernel-driver/drivers/gcip
 
-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 $(GCIP_DIR)/gcip.o
+edgetpu-objs	:= edgetpu-async.o edgetpu-dmabuf.o edgetpu-iremap-pool.o \
+		   edgetpu-kci.o edgetpu-mailbox.o edgetpu-mapping.o \
+		   edgetpu-sw-watchdog.o edgetpu-telemetry.o \
+		   edgetpu-firmware-util.o edgetpu-firmware.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)
+ifndef CONFIG_EDGETPU_TEST
+# Unit testing doesn't need this because GCIP is compiled as a built-in there.
+edgetpu-objs	+= $(GCIP_DIR)/gcip.o
+endif
+
+rio-objs	:= rio-core.o rio-debug-dump.o rio-device-group.o rio-device.o \
+		   rio-firmware.o rio-fs.o rio-iommu.o rio-platform.o rio-pm.o \
+		   rio-thermal.o rio-usage-stats.o rio-wakelock.o rio-external.o \
+		   rio-soc.o $(edgetpu-objs)
 
 CFLAGS_rio-fs.o := -DCONFIG_RIO=1
 CFLAGS_rio-core.o := -DCONFIG_RIO=1
@@ -33,3 +49,4 @@ 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 177bbb5..09895a1 100644
--- a/drivers/edgetpu/Makefile
+++ b/drivers/edgetpu/Makefile
@@ -3,45 +3,20 @@
 # Makefile for EdgeTPU framework and dependent drivers.
 #
 
-obj-$(CONFIG_RIO)		+= rio.o
-
 KERNEL_SRC ?= /lib/modules/$(shell uname -r)/build
 M ?= $(shell pwd)
 
-GIT_PATH=$(KERNEL_SRC)/../private/google-modules/edgetpu/rio
-ifeq ($(shell git --git-dir=$(GIT_PATH)/.git rev-parse --is-inside-work-tree),true)
-	GIT_REPO_STATE=$(shell (git --git-dir=$(GIT_PATH)/.git --work-tree=$(GIT_PATH) status --porcelain | grep -q .) && echo -dirty)
-	ccflags-y       += -DGIT_REPO_TAG=\"$(shell git --git-dir=$(GIT_PATH)/.git rev-parse --short HEAD)$(GIT_REPO_STATE)\"
-else
-	ccflags-y       += -DGIT_REPO_TAG=\"Not\ a\ git\ repository\"
-endif
-
-edgetpu-objs	:= edgetpu-async.o edgetpu-dmabuf.o edgetpu-iremap-pool.o \
-		   edgetpu-kci.o edgetpu-mailbox.o edgetpu-mapping.o \
-		   edgetpu-sw-watchdog.o edgetpu-telemetry.o \
-		   edgetpu-firmware-util.o edgetpu-firmware.o
-
-rio-objs	:= rio-core.o rio-debug-dump.o rio-device-group.o rio-device.o \
-		   rio-firmware.o rio-fs.o rio-iommu.o rio-platform.o rio-pm.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
+KBUILD_OPTIONS	+= CONFIG_GCIP=y
 
-GCIP_DIR=gcip-kernel-driver/drivers/gcip
+-include $(KERNEL_SRC)/../private/google-modules/soc/gs/Makefile.include
 
-KERNEL_SRC ?= /lib/modules/$(shell uname -r)/build
-M ?= $(shell pwd)
+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)" $(@)
+	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)" $(@)
+	$(MAKE) -C $(KERNEL_SRC) M=$(M)/$(GCIP_DIR) W=1 $(@)
+	$(MAKE) -C $(KERNEL_SRC) M=$(M) W=1 $(KBUILD_OPTIONS) $(@)
diff --git a/drivers/edgetpu/edgetpu-device-group.c b/drivers/edgetpu/edgetpu-device-group.c
index 91e83d7..45bb040 100644
--- a/drivers/edgetpu/edgetpu-device-group.c
+++ b/drivers/edgetpu/edgetpu-device-group.c
@@ -27,6 +27,7 @@
 #include "edgetpu-async.h"
 #include "edgetpu-config.h"
 #include "edgetpu-device-group.h"
+#include "edgetpu-dmabuf.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-kci.h"
@@ -280,6 +281,8 @@ static void edgetpu_device_group_release(struct edgetpu_device_group *group)
 		edgetpu_mmu_detach_domain(group->etdev, group->etdomain);
 		edgetpu_mmu_free_domain(group->etdev, group->etdomain);
 	}
+	/* Signal any unsignaled dma fences owned by the group with an error. */
+	edgetpu_sync_fence_group_shutdown(group);
 	group->status = EDGETPU_DEVICE_GROUP_DISBANDED;
 }
 
@@ -451,6 +454,7 @@ edgetpu_device_group_alloc(struct edgetpu_client *client,
 	group->vii.etdev = client->etdev;
 	mutex_init(&group->lock);
 	rwlock_init(&group->events.lock);
+	INIT_LIST_HEAD(&group->dma_fence_list);
 	edgetpu_mapping_init(&group->host_mappings);
 	edgetpu_mapping_init(&group->dmabuf_mappings);
 	group->mbox_attr = *attr;
diff --git a/drivers/edgetpu/edgetpu-device-group.h b/drivers/edgetpu/edgetpu-device-group.h
index 0abf3a2..57fe6e5 100644
--- a/drivers/edgetpu/edgetpu-device-group.h
+++ b/drivers/edgetpu/edgetpu-device-group.h
@@ -109,6 +109,9 @@ struct edgetpu_device_group {
 	/* Mask of errors set for this group. */
 	uint fatal_errors;
 
+	/* List of DMA fences owned by this group */
+	struct list_head dma_fence_list;
+
 	/* end of fields protected by @lock */
 
 	/* TPU IOVA mapped to host DRAM space */
diff --git a/drivers/edgetpu/edgetpu-dmabuf.c b/drivers/edgetpu/edgetpu-dmabuf.c
index c139cda..2b794f1 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.c
+++ b/drivers/edgetpu/edgetpu-dmabuf.c
@@ -60,6 +60,9 @@ struct edgetpu_dmabuf_map {
  * @fence:		the base DMA fence
  * @lock:		spinlock protecting updates to @fence
  * @timeline_name:	name of the timeline associated with the fence
+ * @group:		owning device group
+ * @etfence_list:	global list of all edgetpu DMA fences
+ * @group_list:		list of DMA fences owned by the same group
  *
  * It is likely timelines will become a separate object in the future,
  * but for now there's a unique named timeline associated with each fence.
@@ -68,7 +71,9 @@ struct edgetpu_dma_fence {
 	struct dma_fence fence;
 	spinlock_t lock;
 	char timeline_name[EDGETPU_SYNC_TIMELINE_NAME_LEN];
+	struct edgetpu_device_group *group;
 	struct list_head etfence_list;
+	struct list_head group_list;
 };
 
 /* List of all edgetpu fence objects for debugging. */
@@ -480,6 +485,7 @@ static const char *edgetpu_dma_fence_get_timeline_name(struct dma_fence *fence)
 static void edgetpu_dma_fence_release(struct dma_fence *fence)
 {
 	struct edgetpu_dma_fence *etfence = to_etfence(fence);
+	struct edgetpu_device_group *group;
 	unsigned long flags;
 
 	if (!etfence)
@@ -488,6 +494,17 @@ static void edgetpu_dma_fence_release(struct dma_fence *fence)
 	spin_lock_irqsave(&etfence_list_lock, flags);
 	list_del(&etfence->etfence_list);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
+
+	/* TODO(b/258868303): Don't remove this check when group required, might not yet be set. */
+	group = etfence->group;
+	if (group) {
+		mutex_lock(&group->lock);
+		list_del(&etfence->group_list);
+		mutex_unlock(&group->lock);
+		/* Release this fence's reference on the owning group. */
+		edgetpu_device_group_put(group);
+	}
+
 	kfree(etfence);
 }
 
@@ -511,7 +528,8 @@ static const struct dma_fence_ops edgetpu_dma_fence_ops = {
 #define SEQ_FMT "%llu"
 #endif
 
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap)
 {
 	int fd = get_unused_fd_flags(O_CLOEXEC);
 	int ret;
@@ -533,6 +551,7 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	 * list_head is needed for list_del().
 	 */
 	INIT_LIST_HEAD(&etfence->etfence_list);
+	INIT_LIST_HEAD(&etfence->group_list);
 	memcpy(&etfence->timeline_name, &datap->timeline_name,
 	       EDGETPU_SYNC_TIMELINE_NAME_LEN - 1);
 
@@ -552,6 +571,14 @@ int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap)
 	list_add_tail(&etfence->etfence_list, &etfence_list_head);
 	spin_unlock_irqrestore(&etfence_list_lock, flags);
 
+	/* TODO(b/258868303): Make group required, disallow creating fence we can't track. */
+	if (group) {
+		etfence->group = edgetpu_device_group_get(group);
+		mutex_lock(&group->lock);
+		list_add_tail(&etfence->group_list, &group->dma_fence_list);
+		mutex_unlock(&group->lock);
+	}
+
 	fd_install(fd, sync_file->file);
 	datap->fence = fd;
 	return 0;
@@ -561,26 +588,14 @@ err_put_fd:
 	return ret;
 }
 
-int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+static int _edgetpu_sync_fence_signal(struct dma_fence *fence, int errno, bool ignore_signaled)
 {
-	struct dma_fence *fence;
-	int errno;
 	int ret;
 
-	errno = datap->error;
-	if (errno > 0)
-		errno = -errno;
-	if (errno < -MAX_ERRNO)
-		return -EINVAL;
-
-	fence = sync_file_get_fence(datap->fence);
-	if (!fence)
-		return -EINVAL;
-
 	spin_lock_irq(fence->lock);
 	/* don't signal fence twice */
 	if (unlikely(test_bit(DMA_FENCE_FLAG_SIGNALED_BIT, &fence->flags))) {
-		ret = -EINVAL;
+		ret = ignore_signaled ? 0 : -EINVAL;
 		goto out_unlock;
 	}
 	pr_debug("%s: %s-%s%llu-" SEQ_FMT " errno=%d\n", __func__,
@@ -593,10 +608,51 @@ int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
 
 out_unlock:
 	spin_unlock_irq(fence->lock);
+	return ret;
+}
+
+int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap)
+{
+	struct dma_fence *fence;
+	int errno;
+	int ret;
+
+	errno = datap->error;
+	if (errno > 0)
+		errno = -errno;
+	if (errno < -MAX_ERRNO)
+		return -EINVAL;
+
+	fence = sync_file_get_fence(datap->fence);
+	if (!fence)
+		return -EINVAL;
+
+	ret = _edgetpu_sync_fence_signal(fence, errno, false);
 	dma_fence_put(fence);
 	return ret;
 }
 
+/* Caller holds group lock. */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group)
+{
+	struct list_head *pos;
+	int ret;
+
+	lockdep_assert_held(&group->lock);
+	list_for_each(pos, &group->dma_fence_list) {
+		struct edgetpu_dma_fence *etfence =
+			container_of(pos, struct edgetpu_dma_fence, group_list);
+		struct dma_fence *fence = &etfence->fence;
+
+		ret = _edgetpu_sync_fence_signal(fence, -EPIPE, true);
+		if (ret)
+			etdev_warn(group->etdev, "error %d signaling fence %s-%s %llu-" SEQ_FMT,
+				   ret, fence->ops->get_driver_name(fence),
+				   fence->ops->get_timeline_name(fence),
+				   fence->context, fence->seqno);
+	}
+}
+
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap)
 {
 	struct dma_fence *fence;
@@ -649,7 +705,9 @@ int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused)
 
 		if (fence->error)
 			seq_printf(s, " err=%d", fence->error);
-
+		/* TODO(b/258868303): Remove check when group is required. */
+		if (etfence->group)
+			seq_printf(s, " group=%u", etfence->group->workload_id);
 		seq_putc(s, '\n');
 		spin_unlock_irq(&etfence->lock);
 	}
diff --git a/drivers/edgetpu/edgetpu-dmabuf.h b/drivers/edgetpu/edgetpu-dmabuf.h
index 50d9a60..70af6fd 100644
--- a/drivers/edgetpu/edgetpu-dmabuf.h
+++ b/drivers/edgetpu/edgetpu-dmabuf.h
@@ -26,11 +26,17 @@ int edgetpu_map_dmabuf(struct edgetpu_device_group *group,
 int edgetpu_unmap_dmabuf(struct edgetpu_device_group *group,
 			 tpu_addr_t tpu_addr);
 /* Create a DMA sync fence via ioctl */
-int edgetpu_sync_fence_create(struct edgetpu_create_sync_fence_data *datap);
+int edgetpu_sync_fence_create(struct edgetpu_device_group *group,
+			      struct edgetpu_create_sync_fence_data *datap);
 /* Signal a DMA sync fence, optionally specifying error status */
 int edgetpu_sync_fence_signal(struct edgetpu_signal_sync_fence_data *datap);
 /* Return DMA sync fence status */
 int edgetpu_sync_fence_status(struct edgetpu_sync_fence_status *datap);
+/*
+ * Send error signal to any remaining unsignalled DMA sync fences in a group being disbanded.
+ * Caller holds group lock.
+ */
+void edgetpu_sync_fence_group_shutdown(struct edgetpu_device_group *group);
 /* Dump sync fence info from debugfs */
 int edgetpu_sync_fence_debugfs_show(struct seq_file *s, void *unused);
 
diff --git a/drivers/edgetpu/edgetpu-external.c b/drivers/edgetpu/edgetpu-external.c
index 89c2278..e0c39d6 100644
--- a/drivers/edgetpu/edgetpu-external.c
+++ b/drivers/edgetpu/edgetpu-external.c
@@ -61,12 +61,20 @@ static int edgetpu_external_mailbox_alloc(struct device *edgetpu_dev,
 	struct edgetpu_external_mailbox *ext_mailbox;
 	struct edgetpu_external_mailbox_req req;
 	int ret = 0;
-	struct fd f = fdget(client_info->tpu_fd);
-	struct file *file = f.file;
+	struct file *file;
+	bool use_file = client_info->tpu_fd == -1;
+
+	if (use_file)
+		file = client_info->tpu_file;
+	else
+		file = fget(client_info->tpu_fd);
 
 	if (!file)
 		return -EBADF;
 
+	if (use_file)
+		get_file(file);
+
 	if (!is_edgetpu_file(file)) {
 		ret = -EINVAL;
 		goto out;
@@ -116,7 +124,7 @@ unlock:
 error_put_group:
 	edgetpu_device_group_put(group);
 out:
-	fdput(f);
+	fput(file);
 	return ret;
 }
 
@@ -125,12 +133,20 @@ static int edgetpu_external_mailbox_free(struct device *edgetpu_dev,
 {
 	struct edgetpu_client *client;
 	int ret = 0;
-	struct fd f = fdget(client_info->tpu_fd);
-	struct file *file = f.file;
+	struct file *file;
+	bool use_file = client_info->tpu_fd == -1;
+
+	if (use_file)
+		file = client_info->tpu_file;
+	else
+		file = fget(client_info->tpu_fd);
 
 	if (!file)
 		return -EBADF;
 
+	if (use_file)
+		get_file(file);
+
 	if (!is_edgetpu_file(file)) {
 		ret = -EINVAL;
 		goto out;
@@ -144,12 +160,57 @@ static int edgetpu_external_mailbox_free(struct device *edgetpu_dev,
 
 	ret = edgetpu_mailbox_disable_ext(client, EDGETPU_MAILBOX_ID_USE_ASSOC);
 out:
-	fdput(f);
+	fput(file);
+	return ret;
+}
+
+static int edgetpu_external_start_offload(struct device *edgetpu_dev,
+					  struct edgetpu_ext_client_info *client_info,
+					  struct edgetpu_ext_offload_info *offload_info)
+{
+	struct edgetpu_client *client;
+	struct edgetpu_device_group *group;
+	struct file *file = client_info->tpu_file;
+	int ret = 0;
+
+	if (!file)
+		return -EBADF;
+
+	get_file(file);
+
+	if (!is_edgetpu_file(file)) {
+		ret = -EINVAL;
+		goto out;
+	}
+
+	client = file->private_data;
+	if (!client || client->etdev->dev != edgetpu_dev) {
+		ret = -EINVAL;
+		goto out;
+	}
+
+	mutex_lock(&client->group_lock);
+	if (!client->group) {
+		ret = -EINVAL;
+		mutex_unlock(&client->group_lock);
+		goto out;
+	}
+	group = edgetpu_device_group_get(client->group);
+	mutex_unlock(&client->group_lock);
+
+	mutex_lock(&group->lock);
+	offload_info->client_id = group->vcid;
+	/* TODO(b/259213063): notify the firmware to listen for the TPU offloads. */
+	mutex_unlock(&group->lock);
+
+	edgetpu_device_group_put(group);
+
+out:
+	fput(file);
 	return ret;
 }
 
-int edgetpu_ext_driver_cmd(struct device *edgetpu_dev,
-			   enum edgetpu_ext_client_type client_type,
+int edgetpu_ext_driver_cmd(struct device *edgetpu_dev, enum edgetpu_ext_client_type client_type,
 			   enum edgetpu_ext_commands cmd_id, void *in_data, void *out_data)
 {
 	switch (cmd_id) {
@@ -157,6 +218,8 @@ int edgetpu_ext_driver_cmd(struct device *edgetpu_dev,
 		return edgetpu_external_mailbox_alloc(edgetpu_dev, in_data, out_data, client_type);
 	case FREE_EXTERNAL_MAILBOX:
 		return edgetpu_external_mailbox_free(edgetpu_dev, in_data);
+	case START_OFFLOAD:
+		return edgetpu_external_start_offload(edgetpu_dev, in_data, out_data);
 	default:
 		return -ENOENT;
 	}
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 39730ec..1038303 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -377,6 +377,7 @@ edgetpu_ioctl_unmap_dmabuf(struct edgetpu_client *client,
 }
 
 static int edgetpu_ioctl_sync_fence_create(
+	struct edgetpu_client *client,
 	struct edgetpu_create_sync_fence_data __user *datap)
 {
 	struct edgetpu_create_sync_fence_data data;
@@ -384,7 +385,13 @@ static int edgetpu_ioctl_sync_fence_create(
 
 	if (copy_from_user(&data, (void __user *)datap, sizeof(data)))
 		return -EFAULT;
-	ret = edgetpu_sync_fence_create(&data);
+	LOCK(client);
+	if (!client->group)
+		/* TODO(b/258868303): Require a group, disallow creating a fence we can't track. */
+		etdev_warn(client->etdev,
+			   "client creating sync fence not joined to a device group");
+	ret = edgetpu_sync_fence_create(client->group, &data);
+	UNLOCK(client);
 	if (ret)
 		return ret;
 	if (copy_to_user((void __user *)datap, &data, sizeof(data)))
@@ -687,7 +694,7 @@ long edgetpu_ioctl(struct file *file, uint cmd, ulong arg)
 		ret = edgetpu_ioctl_allocate_device_buffer(client, (u64)argp);
 		break;
 	case EDGETPU_CREATE_SYNC_FENCE:
-		ret = edgetpu_ioctl_sync_fence_create(argp);
+		ret = edgetpu_ioctl_sync_fence_create(client, argp);
 		break;
 	case EDGETPU_SIGNAL_SYNC_FENCE:
 		ret = edgetpu_ioctl_sync_fence_signal(argp);
diff --git a/drivers/edgetpu/include/linux/gsa/gsa_tpu.h b/drivers/edgetpu/include/linux/gsa/gsa_tpu.h
deleted file mode 100644
index 2dc8a79..0000000
--- a/drivers/edgetpu/include/linux/gsa/gsa_tpu.h
+++ /dev/null
@@ -1,53 +0,0 @@
-/* SPDX-License-Identifier: GPL-2.0 */
-/*
- * Fallback header for systems without GSA support.
- *
- * Copyright (C) 2021 Google, Inc.
- */
-
-#ifndef __LINUX_GSA_TPU_H
-#define __LINUX_GSA_TPU_H
-
-#include <linux/device.h>
-#include <linux/types.h>
-
-#if IS_ENABLED(CONFIG_EDGETPU_TEST)
-
-int gsa_load_tpu_fw_image(struct device *gsa, dma_addr_t img_meta, phys_addr_t img_body);
-
-#else /* IS_ENABLED(CONFIG_EDGETPU_TEST) */
-
-static inline int gsa_load_tpu_fw_image(struct device *gsa, dma_addr_t img_meta,
-					phys_addr_t img_body)
-{
-	return 0;
-}
-
-#endif /* IS_ENABLED(CONFIG_EDGETPU_TEST) */
-
-static inline int gsa_unload_tpu_fw_image(struct device *gsa)
-{
-	return 0;
-}
-
-enum gsa_tpu_state {
-	GSA_TPU_STATE_INACTIVE = 0,
-	GSA_TPU_STATE_LOADED,
-	GSA_TPU_STATE_RUNNING,
-	GSA_TPU_STATE_SUSPENDED,
-};
-
-enum gsa_tpu_cmd {
-	GSA_TPU_GET_STATE = 0,
-	GSA_TPU_START,
-	GSA_TPU_SUSPEND,
-	GSA_TPU_RESUME,
-	GSA_TPU_SHUTDOWN,
-};
-
-static inline int gsa_send_tpu_cmd(struct device *gsa, enum gsa_tpu_cmd cmd)
-{
-	return 0;
-}
-
-#endif /* __LINUX_GSA_TPU_H */
diff --git a/drivers/edgetpu/include/soc/google/gs_tmu.h b/drivers/edgetpu/include/soc/google/gs_tmu.h
deleted file mode 100644
index 2c74634..0000000
--- a/drivers/edgetpu/include/soc/google/gs_tmu.h
+++ /dev/null
@@ -1,25 +0,0 @@
-/* SPDX-License-Identifier: GPL-2.0-only
- *
- * gs_tmu.h - defines and stubs to build edgetpu driver unit tests
- *
- * Based on gs101_tmu.h
- *  Copyright (C) 2019 Samsung Electronics
- *  Hyeonseong Gil <hs.gill@samsung.com>
- */
-
-#ifndef _GS_TMU_H
-#define _GS_TMU_H
-
-enum thermal_pause_state {
-	THERMAL_RESUME = 0,
-	THERMAL_SUSPEND,
-};
-
-typedef int (*tpu_pause_cb)(enum thermal_pause_state action, void *data);
-
-static inline void
-register_tpu_thermal_pause_cb(tpu_pause_cb tpu_cb, void *data)
-{
-}
-
-#endif /* _GS_TMU_H */
diff --git a/drivers/edgetpu/include/trace/events/edgetpu.h b/drivers/edgetpu/include/trace/events/edgetpu.h
index 9723eca..e8b9e5d 100644
--- a/drivers/edgetpu/include/trace/events/edgetpu.h
+++ b/drivers/edgetpu/include/trace/events/edgetpu.h
@@ -128,4 +128,8 @@ TRACE_EVENT(edgetpu_map_dmabuf_end,
 #endif /* _TRACE_EDGETPU_H */
 
 /* This part must be outside protection */
+#undef TRACE_INCLUDE_PATH
+#define TRACE_INCLUDE_PATH .
+#define TRACE_INCLUDE_FILE trace/events/edgetpu
+
 #include <trace/define_trace.h>
diff --git a/drivers/edgetpu/mobile-soc-gsx01.c b/drivers/edgetpu/mobile-soc-gsx01.c
index baefb33..8fb5288 100644
--- a/drivers/edgetpu/mobile-soc-gsx01.c
+++ b/drivers/edgetpu/mobile-soc-gsx01.c
@@ -14,12 +14,13 @@
 #include <linux/types.h>
 #include <soc/google/bts.h>
 #include <soc/google/exynos_pm_qos.h>
-#include <soc/google/gs_tmu.h>
+#include <soc/google/gs_tmu_v3.h>
 
 #include "edgetpu-internal.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mobile-platform.h"
+#include "edgetpu-pm.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-thermal.h"
 #include "mobile-firmware.h"
@@ -302,8 +303,6 @@ static unsigned long edgetpu_pm_rate;
 
 long edgetpu_soc_pm_get_rate(struct edgetpu_dev *etdev, int flags)
 {
-	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
-	struct edgetpu_mobile_platform_pwr *platform_pwr = &etmdev->platform_pwr;
 	void __iomem *cmu_base = etdev->soc_data->cmu_base;
 	long curr_state;
 	u32 pll_con3;
@@ -314,13 +313,11 @@ long edgetpu_soc_pm_get_rate(struct edgetpu_dev *etdev, int flags)
 	if (!cmu_base)
 		return -EINVAL;
 
-	if (!platform_pwr->is_block_down)
-		etdev_warn(etdev,
-			   "Querying the CMU PLL register when TPU_OFF might lead to crash.");
-	else if (platform_pwr->is_block_down(etdev))
+	/* We need to keep TPU being powered to ensure CMU read is valid. */
+	if (!edgetpu_pm_get_if_powered(etdev->pm))
 		return 0;
-
 	pll_con3 = readl(cmu_base + PLL_CON3_OFFSET);
+	edgetpu_pm_put(etdev->pm);
 
 	/*
 	 * Below values must match the CMU PLL (pll_con3_pll_tpu) values in the spec and firmware.