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

gcip: implement noncontiguous alloc
Bug: 262684159
gcip: add gcip-alloc-helper.h
Bug: 262684159 (repeat)
edgetpu: return context_id as client_id to DSP
Bug: 262419733
gcip: Change hard-coded magic numbers to MACROs
Bug: 257300340
edgetpu: Add wakelock trace
Bug: 262914320
edgetpu: Fix missing wakelock unlock
gcip: Update the comments in gcip-image-config for new encoding
Bug: 257300340 (repeat)

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

8 files changed, 279 insertions, 38 deletions

diff --git a/drivers/edgetpu/edgetpu-external.c b/drivers/edgetpu/edgetpu-external.c
index e0c39d6..adf8d3b 100644
--- a/drivers/edgetpu/edgetpu-external.c
+++ b/drivers/edgetpu/edgetpu-external.c
@@ -171,6 +171,7 @@ static int edgetpu_external_start_offload(struct device *edgetpu_dev,
 	struct edgetpu_client *client;
 	struct edgetpu_device_group *group;
 	struct file *file = client_info->tpu_file;
+	enum edgetpu_context_id context;
 	int ret = 0;
 
 	if (!file)
@@ -199,12 +200,17 @@ static int edgetpu_external_start_offload(struct device *edgetpu_dev,
 	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);
+	context = edgetpu_group_context_id_locked(group);
+	if (context == EDGETPU_CONTEXT_INVALID || context & EDGETPU_CONTEXT_DOMAIN_TOKEN) {
+		ret = -EINVAL;
+		goto out_group_unlock;
+	}
 
-	edgetpu_device_group_put(group);
+	offload_info->client_id = context;
 
+out_group_unlock:
+	mutex_unlock(&group->lock);
+	edgetpu_device_group_put(group);
 out:
 	fput(file);
 	return ret;
diff --git a/drivers/edgetpu/edgetpu-fs.c b/drivers/edgetpu/edgetpu-fs.c
index 1038303..03b07f4 100644
--- a/drivers/edgetpu/edgetpu-fs.c
+++ b/drivers/edgetpu/edgetpu-fs.c
@@ -467,12 +467,17 @@ static int edgetpu_ioctl_release_wakelock(struct edgetpu_client *client)
 {
 	int count;
 
+	trace_edgetpu_release_wakelock_start(client->pid);
+
 	LOCK(client);
 	edgetpu_wakelock_lock(client->wakelock);
 	count = edgetpu_wakelock_release(client->wakelock);
 	if (count < 0) {
 		edgetpu_wakelock_unlock(client->wakelock);
 		UNLOCK(client);
+
+		trace_edgetpu_release_wakelock_end(client->pid, count);
+
 		return count;
 	}
 	if (!count) {
@@ -484,15 +489,20 @@ static int edgetpu_ioctl_release_wakelock(struct edgetpu_client *client)
 	UNLOCK(client);
 	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
 		  count);
+
+	trace_edgetpu_release_wakelock_end(client->pid, count);
+
 	return 0;
 }
 
 static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 {
-	int count;
+	int count = 0;
 	int ret;
 	struct edgetpu_thermal *thermal = client->etdev->thermal;
 
+	trace_edgetpu_acquire_wakelock_start(current->pid);
+
 	LOCK(client);
 	/*
 	 * Update client PID; the client may have been passed from the
@@ -509,46 +519,48 @@ static int edgetpu_ioctl_acquire_wakelock(struct edgetpu_client *client)
 		etdev_warn_ratelimited(client->etdev,
 				       "wakelock acquire rejected due to thermal suspend");
 		edgetpu_thermal_unlock(thermal);
-		goto error_unlock;
+		goto error_client_unlock;
 	} else {
 		ret = edgetpu_pm_get(client->etdev->pm);
 		edgetpu_thermal_unlock(thermal);
 		if (ret) {
 			etdev_warn(client->etdev, "pm_get failed (%d)", ret);
-			goto error_unlock;
+			goto error_client_unlock;
 		}
 	}
 	edgetpu_wakelock_lock(client->wakelock);
 	count = edgetpu_wakelock_acquire(client->wakelock);
 	if (count < 0) {
-		edgetpu_pm_put(client->etdev->pm);
 		ret = count;
-		goto error_unlock;
+		goto error_wakelock_unlock;
 	}
 	if (!count) {
 		if (client->group)
 			ret = edgetpu_group_attach_and_open_mailbox(client->group);
 		if (ret) {
-			etdev_warn(client->etdev,
-				   "failed to attach mailbox: %d", ret);
-			edgetpu_pm_put(client->etdev->pm);
+			etdev_warn(client->etdev, "failed to attach mailbox: %d", ret);
 			edgetpu_wakelock_release(client->wakelock);
-			edgetpu_wakelock_unlock(client->wakelock);
-			goto error_unlock;
+			goto error_wakelock_unlock;
 		}
-	} else {
-		/* Balance the power up count due to pm_get above.*/
-		edgetpu_pm_put(client->etdev->pm);
 	}
+
+error_wakelock_unlock:
 	edgetpu_wakelock_unlock(client->wakelock);
+
+	/* Balance the power up count due to pm_get above.*/
+	if (ret || count)
+		edgetpu_pm_put(client->etdev->pm);
+
+error_client_unlock:
 	UNLOCK(client);
-	etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__,
-		  count + 1);
-	return 0;
-error_unlock:
-	UNLOCK(client);
-	etdev_err(client->etdev, "client pid %d failed to acquire wakelock",
-		  client->pid);
+
+	if (ret)
+		etdev_err(client->etdev, "client pid %d failed to acquire wakelock", client->pid);
+	else
+		etdev_dbg(client->etdev, "%s: wakelock req count = %u", __func__, count + 1);
+
+	trace_edgetpu_acquire_wakelock_end(client->pid, ret ? count : count + 1, ret);
+
 	return ret;
 }
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
index 2f34448..c3424ee 100644
--- a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/Makefile
@@ -6,8 +6,9 @@
 CONFIG_GCIP		?= m
 obj-$(CONFIG_GCIP)	+= gcip.o
 
-gcip-objs		:= gcip-domain-pool.o gcip-firmware.o gcip-image-config.o gcip-kci.o \
-			   gcip-mailbox.o gcip-mem-pool.o gcip-telemetry.o
+gcip-objs		:= gcip-alloc-helper.o gcip-domain-pool.o gcip-firmware.o \
+			   gcip-image-config.o gcip-kci.o gcip-mailbox.o gcip-mem-pool.o \
+			   gcip-telemetry.o
 
 CURRENT_DIR=$(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 
diff --git a/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-alloc-helper.c b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-alloc-helper.c
new file mode 100644
index 0000000..f79149f
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/drivers/gcip/gcip-alloc-helper.c
@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * GCIP helpers for allocating memories.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#include <asm/page.h>
+#include <linux/device.h>
+#include <linux/mm_types.h>
+#include <linux/scatterlist.h>
+#include <linux/slab.h>
+#include <linux/vmalloc.h>
+
+#include <gcip/gcip-alloc-helper.h>
+
+/*
+ * Set @pages to the pages @mem represents.
+ * @mem must be a pointer returned by vmalloc.
+ *
+ * Returns 0 on success, -ENOMEM when any page is NULL.
+ */
+static int gcip_vmalloc_to_pages(void *mem, size_t count, struct page **pages)
+{
+	size_t i = 0;
+
+	while (count--) {
+		pages[i++] = vmalloc_to_page(mem);
+		if (!pages[i - 1])
+			return -ENOMEM;
+		mem += PAGE_SIZE;
+	}
+	return 0;
+}
+
+struct sg_table *gcip_alloc_noncontiguous(struct device *dev, size_t size, gfp_t gfp)
+{
+	struct gcip_sgt_handle *sh = kmalloc(sizeof(*sh), gfp);
+	void *mem;
+	struct page **pages;
+	size_t count;
+	int ret;
+
+	if (!sh)
+		return NULL;
+
+	size = PAGE_ALIGN(size);
+	count = size >> PAGE_SHIFT;
+	mem = vzalloc_node(size, dev_to_node(dev));
+	if (!mem) {
+		dev_err(dev, "GCIP noncontiguous alloc size=%#zx failed", size);
+		goto err_free_sh;
+	}
+
+	pages = kmalloc_array(count, sizeof(*pages), gfp);
+	if (!pages) {
+		dev_err(dev, "GCIP alloc pages array count=%zu failed", count);
+		goto err_free_mem;
+	}
+
+	if (gcip_vmalloc_to_pages(mem, count, pages)) {
+		dev_err(dev, "convert memory to pages failed");
+		goto err_free_pages;
+	}
+
+	ret = sg_alloc_table_from_pages(&sh->sgt, pages, count, 0, size, gfp);
+	if (ret) {
+		dev_err(dev, "alloc SG table with size=%#zx failed: %d", size, ret);
+		goto err_free_pages;
+	}
+
+	kfree(pages);
+	sh->mem = mem;
+	return &sh->sgt;
+
+err_free_pages:
+	kfree(pages);
+err_free_mem:
+	vfree(mem);
+err_free_sh:
+	kfree(sh);
+	return NULL;
+}
+
+void gcip_free_noncontiguous(struct sg_table *sgt)
+{
+	struct gcip_sgt_handle *sh = container_of(sgt, struct gcip_sgt_handle, sgt);
+
+	sg_free_table(&sh->sgt);
+	vfree(sh->mem);
+	kfree(sh);
+}
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 a440765..5fed69c 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
@@ -12,21 +12,26 @@
 
 #include <gcip/gcip-image-config.h>
 
-#define SIZE_MASK 0xfff
+#define ADDR_SHIFT 12
+#define SIZE_MODE_BIT BIT(ADDR_SHIFT - 1)
+#define SECURE_SIZE_MASK (SIZE_MODE_BIT - 1u)
+#define NS_SIZE_MASK (BIT(ADDR_SHIFT) - 1u)
+#define ADDR_MASK ~(BIT(ADDR_SHIFT) - 1u)
 
 /* used by ns_iommu_mappings */
-#define CONFIG_TO_MBSIZE(a) (((a)&SIZE_MASK) << 20)
+#define CONFIG_TO_MBSIZE(a) (((a)&NS_SIZE_MASK) << 20)
 
 /* used by iommu_mappings */
 static inline __u32 config_to_size(__u32 cfg)
 {
 	__u32 page_size;
 
-	if (cfg & 0x800)
-		page_size = cfg & 0x7ff;
+	if (cfg & SIZE_MODE_BIT)
+		page_size = cfg & SECURE_SIZE_MASK;
 	else
-		page_size = 1U << (cfg & SIZE_MASK);
-	return page_size << 12;
+		page_size = BIT(cfg & SECURE_SIZE_MASK);
+
+	return page_size << PAGE_SHIFT;
 }
 
 static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
@@ -45,7 +50,7 @@ static int setup_iommu_mappings(struct gcip_image_config_parser *parser,
 			goto err;
 		}
 		size = config_to_size(config->iommu_mappings[i].image_config_value);
-		paddr = config->iommu_mappings[i].image_config_value & ~SIZE_MASK;
+		paddr = config->iommu_mappings[i].image_config_value & ADDR_MASK;
 
 		dev_dbg(parser->dev, "Image config adding IOMMU mapping: %pad -> %pap", &daddr,
 			&paddr);
@@ -100,7 +105,7 @@ static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 	phys_addr_t paddr = 0;
 
 	for (i = 0; i < config->num_ns_iommu_mappings; i++) {
-		daddr = config->ns_iommu_mappings[i] & ~SIZE_MASK;
+		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
 		if (unlikely(!daddr)) {
 			dev_warn(parser->dev, "Invalid config, device address is zero");
 			ret = -EIO;
@@ -124,7 +129,7 @@ static int setup_ns_iommu_mappings(struct gcip_image_config_parser *parser,
 err:
 	while (i--) {
 		size = CONFIG_TO_MBSIZE(config->ns_iommu_mappings[i]);
-		daddr = config->ns_iommu_mappings[i] & ~SIZE_MASK;
+		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
 		parser->ops->unmap(parser->data, daddr, size, 0);
 	}
 	return ret;
@@ -139,7 +144,7 @@ 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;
+		daddr = config->ns_iommu_mappings[i] & ADDR_MASK;
 		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/gcip-kernel-driver/include/gcip/gcip-alloc-helper.h b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-alloc-helper.h
new file mode 100644
index 0000000..3d2c110
--- /dev/null
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-alloc-helper.h
@@ -0,0 +1,50 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+/*
+ * GCIP helpers for allocating memories.
+ *
+ * Copyright (C) 2022 Google LLC
+ */
+
+#ifndef __GCIP_ALLOC_HELPER_H__
+#define __GCIP_ALLOC_HELPER_H__
+
+#include <linux/device.h>
+#include <linux/scatterlist.h>
+#include <linux/types.h>
+
+/*
+ * The actual return value from the alloc_noncontiguous function.
+ * The user should only care about @sgt. @pages is used internally for freeing memory.
+ */
+struct gcip_sgt_handle {
+	struct sg_table sgt;
+	void *mem;
+};
+
+/*
+ * Allocates non-contiguous memory with size @size bytes.
+ *
+ * @dev: pointer to device structure. Is used for logging or the NUMA node for page allocation.
+ * @size: Total size in bytes. Will be page aligned.
+ * @gfp: The GFP flag for malloc internal structures.
+ *
+ * Returns the SG table represents the non-contiguous region.
+ * Returns NULL on any error.
+ */
+struct sg_table *gcip_alloc_noncontiguous(struct device *dev, size_t size, gfp_t gfp);
+/* Frees the memory allocated by gcip_alloc_noncontiguous. */
+void gcip_free_noncontiguous(struct sg_table *sgt);
+
+/*
+ * Returns the virtual memory that was used to allocate @sgt.
+ *
+ * @sgt must be the return pointer of gcip_alloc_noncontiguous.
+ */
+static inline void *gcip_noncontiguous_sgt_to_mem(struct sg_table *sgt)
+{
+	struct gcip_sgt_handle *sh = container_of(sgt, struct gcip_sgt_handle, sgt);
+
+	return sh->mem;
+}
+
+#endif /* __GCIP_ALLOC_HELPER_H__ */
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
index a3539a0..a995188 100644
--- a/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
+++ b/drivers/edgetpu/gcip-kernel-driver/include/gcip/gcip-image-config.h
@@ -35,8 +35,9 @@ struct gcip_image_config {
 		/* 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.
+		 * Encodes a 12-bit aligned address and the corresponding size
+		 * into a 32-bit value.
+		 * Detailed encoding method is defined in gcip-image-config.c.
 		 */
 		__u32 image_config_value;
 	} iommu_mappings[GCIP_IMG_CFG_MAX_IOMMU_MAPPINGS];
diff --git a/drivers/edgetpu/include/trace/events/edgetpu.h b/drivers/edgetpu/include/trace/events/edgetpu.h
index e8b9e5d..53669f9 100644
--- a/drivers/edgetpu/include/trace/events/edgetpu.h
+++ b/drivers/edgetpu/include/trace/events/edgetpu.h
@@ -125,6 +125,80 @@ TRACE_EVENT(edgetpu_map_dmabuf_end,
 		__entry->dmabuf_fd, __entry->flags, __entry->die_index)
 );
 
+TRACE_EVENT(edgetpu_acquire_wakelock_start,
+
+	TP_PROTO(pid_t pid),
+
+	TP_ARGS(pid),
+
+	TP_STRUCT__entry(
+		__field(pid_t, pid)
+	),
+
+	TP_fast_assign(
+		__entry->pid = pid;
+	),
+
+	TP_printk("pid = %d", __entry->pid)
+);
+
+TRACE_EVENT(edgetpu_acquire_wakelock_end,
+
+	TP_PROTO(pid_t pid, int count, int ret),
+
+	TP_ARGS(pid, count, ret),
+
+	TP_STRUCT__entry(
+		__field(pid_t, pid)
+		__field(int, count)
+		__field(int, ret)
+	),
+
+	TP_fast_assign(
+		__entry->pid = pid;
+		__entry->count = count;
+		__entry->ret = ret;
+	),
+
+	TP_printk("pid = %d, req_count = %d, ret = %d", __entry->pid, __entry->count, __entry->ret)
+);
+
+TRACE_EVENT(edgetpu_release_wakelock_start,
+
+	TP_PROTO(pid_t pid),
+
+	TP_ARGS(pid),
+
+	TP_STRUCT__entry(
+		__field(pid_t, pid)
+	),
+
+	TP_fast_assign(
+		__entry->pid = pid;
+	),
+
+	TP_printk("pid = %d", __entry->pid)
+);
+
+TRACE_EVENT(edgetpu_release_wakelock_end,
+
+	TP_PROTO(pid_t pid, int count),
+
+	TP_ARGS(pid, count),
+
+	TP_STRUCT__entry(
+		__field(pid_t, pid)
+		__field(int, count)
+	),
+
+	TP_fast_assign(
+		__entry->pid = pid;
+		__entry->count = count;
+	),
+
+	TP_printk("pid = %d, req_count = %d", __entry->pid, __entry->count)
+);
+
 #endif /* _TRACE_EDGETPU_H */
 
 /* This part must be outside protection */