drivers/edgetpu/mobile-firmware.c - kernel/google-modules/edgetpu/rio - Git at Google

 // SPDX-License-Identifier: GPL-2.0
 /*
  * Edge TPU firmware management for mobile chipsets.
  *
  * Copyright (C) 2021-2022 Google LLC
  */

 #include <linux/device.h>
 #include <linux/dma-mapping.h>
 #include <linux/firmware.h>
 #include <linux/gsa/gsa_tpu.h>
 #include <linux/slab.h>
 #include <linux/string.h>
 #include <linux/types.h>

 #include <gcip/gcip-alloc-helper.h>
 #include <gcip/gcip-image-config.h>

 #include "edgetpu.h"
 #include "edgetpu-config.h"
 #include "edgetpu-firmware.h"
 #include "edgetpu-firmware-util.h"
 #include "edgetpu-internal.h"
 #include "edgetpu-iremap-pool.h"
 #include "edgetpu-kci.h"
 #include "edgetpu-mailbox.h"
 #include "edgetpu-mmu.h"
 #include "edgetpu-mobile-platform.h"
 #include "edgetpu-soc.h"
 #include "edgetpu-telemetry.h"
 #include "mobile-firmware.h"

 static int image_config_map(void *data, dma_addr_t daddr, phys_addr_t paddr, size_t size,
 			    unsigned int flags)
 {
 	struct edgetpu_dev *etdev = data;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_fw_ctx *fw_ctx;
 	int ret;

 	if (flags & GCIP_IMAGE_CONFIG_FLAGS_SECURE)
 		return edgetpu_mmu_add_translation(etdev, daddr, paddr, size,
 						   IOMMU_READ | IOMMU_WRITE, EDGETPU_CONTEXT_KCI);

 	fw_ctx = kzalloc(sizeof(*fw_ctx), GFP_KERNEL);
 	if (!fw_ctx)
 		return -ENOMEM;

 	fw_ctx->daddr = daddr;
 	fw_ctx->size = size;
 	fw_ctx->sgt = gcip_alloc_noncontiguous(etdev->dev, size, GFP_KERNEL);
 	if (!fw_ctx->sgt) {
 		kfree(fw_ctx);
 		return -ENOMEM;
 	}

 	ret = edgetpu_mmu_map_iova_sgt(etdev, daddr, fw_ctx->sgt, DMA_BIDIRECTIONAL, 0,
 				       EDGETPU_CONTEXT_KCI);
 	if (ret) {
 		gcip_free_noncontiguous(fw_ctx->sgt);
 		kfree(fw_ctx);
 		return ret;
 	}

 	mutex_lock(&etmdev->fw_ctx_list_lock);
 	list_add_tail(&fw_ctx->list, &etmdev->fw_ctx_list);
 	mutex_unlock(&etmdev->fw_ctx_list_lock);

 	return 0;
 }

 static void image_config_unmap(void *data, dma_addr_t daddr, size_t size, unsigned int flags)
 {
 	struct edgetpu_dev *etdev = data;
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct edgetpu_mobile_fw_ctx *fw_ctx = NULL, *cur;

 	if (flags & GCIP_IMAGE_CONFIG_FLAGS_SECURE) {
 		edgetpu_mmu_remove_translation(etdev, daddr, size, EDGETPU_CONTEXT_KCI);
 		return;
 	}

 	mutex_lock(&etmdev->fw_ctx_list_lock);
 	list_for_each_entry(cur, &etmdev->fw_ctx_list, list) {
 		if (cur->daddr == daddr && cur->size == size) {
 			fw_ctx = cur;
 			list_del(&cur->list);
 			break;
 		}
 	}
 	mutex_unlock(&etmdev->fw_ctx_list_lock);

 	if (fw_ctx) {
 		edgetpu_mmu_unmap_iova_sgt(etdev, daddr, fw_ctx->sgt, DMA_BIDIRECTIONAL,
 					   EDGETPU_CONTEXT_KCI);
 		gcip_free_noncontiguous(fw_ctx->sgt);
 		kfree(fw_ctx);
 	} else {
 		etdev_warn(etdev, "Firmware context region SG table not found.");
 	}
 }

 static int mobile_firmware_after_create(struct edgetpu_firmware *et_fw)
 {
 	struct gcip_image_config_parser *data;
 	struct edgetpu_dev *etdev = et_fw->etdev;
 	int ret;
 	static const struct gcip_image_config_ops image_config_parser_ops = {
 		.map = image_config_map,
 		.unmap = image_config_unmap,
 	};

 	data = kzalloc(sizeof(*data), GFP_KERNEL);
 	if (!data)
 		return -ENOMEM;

 	ret = gcip_image_config_parser_init(data, &image_config_parser_ops,
 					    get_dev_for_logging(etdev), etdev);
 	if (ret) {
 		etdev_err(etdev, "Image config parser init failed: %d", ret);
 		kfree(data);
 		return ret;
 	}

 	edgetpu_firmware_set_data(et_fw, data);
 	return 0;
 }

 static void mobile_firmware_before_destroy(struct edgetpu_firmware *et_fw)
 {
 	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(et_fw);

 	gcip_image_config_clear(cfg_parser);
 	edgetpu_firmware_set_data(et_fw, NULL);
 	kfree(cfg_parser);
 }

 static int mobile_firmware_alloc_buffer(
 		struct edgetpu_firmware *et_fw,
 		struct edgetpu_firmware_buffer *fw_buf)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;

 	/* Allocate extra space the image header */
 	size_t buffer_size = EDGETPU_MAX_FW_LIMIT + MOBILE_FW_HEADER_SIZE;

 	fw_buf->vaddr = vmalloc(buffer_size);
 	if (!fw_buf->vaddr) {
 		etdev_err(etdev, "%s: failed to allocate buffer (%zu bytes)\n",
 			  __func__, buffer_size);
 		return -ENOMEM;
 	}
 	fw_buf->dma_addr = 0;
 	fw_buf->alloc_size = buffer_size;
 	fw_buf->used_size_align = 16;
 	return 0;
 }

 static void mobile_firmware_free_buffer(
 		struct edgetpu_firmware *et_fw,
 		struct edgetpu_firmware_buffer *fw_buf)
 {
 	vfree(fw_buf->vaddr);
 	fw_buf->vaddr = NULL;
 	fw_buf->dma_addr = 0;
 	fw_buf->alloc_size = 0;
 	fw_buf->used_size_align = 0;
 }

 static struct gcip_image_config *mobile_firmware_get_image_config(struct edgetpu_dev *etdev)
 {
 	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(etdev->firmware);

 	return &cfg_parser->last_config;
 }

 static int mobile_firmware_gsa_authenticate(struct edgetpu_mobile_platform_dev *etmdev,
 					    struct edgetpu_firmware_buffer *fw_buf,
 					    void *image_vaddr)
 {
 	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
 	void *header_vaddr;
 	dma_addr_t header_dma_addr;
 	int tpu_state;
 	int ret = 0;

 	tpu_state = gsa_send_tpu_cmd(etmdev->gsa_dev, GSA_TPU_GET_STATE);

 	if (tpu_state < GSA_TPU_STATE_INACTIVE) {
 		etdev_err(etdev, "GSA failed to retrieve current status: %d\n", tpu_state);
 		return tpu_state;
 	}

 	etdev_dbg(etdev, "GSA Reports TPU state: %d\n", tpu_state);

 	if (tpu_state > GSA_TPU_STATE_INACTIVE) {
 		ret = gsa_unload_tpu_fw_image(etmdev->gsa_dev);
 		if (ret) {
 			etdev_warn(etdev, "GSA release failed: %d\n", ret);
 			return -EIO;
 		}
 	}

 	/* Copy the firmware image to the target location, skipping the header */
 	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
 	       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);

 	/* Allocate coherent memory for the image header */
 	header_vaddr = dma_alloc_coherent(etmdev->gsa_dev,
 					  MOBILE_FW_HEADER_SIZE,
 					  &header_dma_addr, GFP_KERNEL);
 	if (!header_vaddr) {
 		etdev_err(etdev,
 			  "Failed to allocate coherent memory for header\n");
 		return -ENOMEM;
 	}

 	memcpy(header_vaddr, fw_buf->vaddr, MOBILE_FW_HEADER_SIZE);
 	etdev_dbg(etdev, "Requesting GSA image load. meta = %llX payload = %llX", header_dma_addr,
 		  (u64)etmdev->fw_region_paddr);

 	ret = gsa_load_tpu_fw_image(etmdev->gsa_dev, header_dma_addr,
 				    etmdev->fw_region_paddr);
 	if (ret)
 		etdev_err(etdev, "GSA authentication failed: %d\n", ret);

 	dma_free_coherent(etmdev->gsa_dev, MOBILE_FW_HEADER_SIZE, header_vaddr, header_dma_addr);

 	return ret;
 }

 static int mobile_firmware_update_remapped_data_region(struct edgetpu_dev *etdev)
 {
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct gcip_image_config *config = mobile_firmware_get_image_config(etdev);
 	tpu_addr_t remapped_data_addr = EDGETPU_INSTRUCTION_REMAP_BASE + etmdev->fw_region_size;
 	size_t remapped_data_size = config->remapped_data_size ? config->remapped_data_size :
 								 EDGETPU_DEFAULT_REMAPPED_DATA_SIZE;
 	int ret;

 	if (etmdev->remapped_data_addr == remapped_data_addr &&
 	    etmdev->remapped_data_size == remapped_data_size)
 		return 0;

 	if (remapped_data_addr < EDGETPU_INSTRUCTION_REMAP_BASE + config->firmware_size ||
 	    config->firmware_base + etmdev->fw_region_size + remapped_data_size >
 		    config->remapped_region_start)
 		return -EINVAL;

 	etdev_dbg(etdev, "Moving remapped data from %#llx to %#llx\n", etmdev->remapped_data_addr,
 		  remapped_data_addr);

 	if (etmdev->shared_mem_vaddr) {
 		/* No need to free the VII queues since allocated groups will block fw loading. */
 		edgetpu_kci_release(etdev, etdev->etkci);
 		edgetpu_telemetry_exit(etdev);
 		edgetpu_iremap_pool_destroy(etdev);
 		memunmap(etmdev->shared_mem_vaddr);
 	}

 	etmdev->remapped_data_addr = remapped_data_addr;
 	etmdev->remapped_data_size = remapped_data_size;

 	etmdev->shared_mem_paddr = etmdev->fw_region_paddr + etmdev->fw_region_size;
 	etmdev->shared_mem_vaddr =
 		memremap(etmdev->shared_mem_paddr, etmdev->remapped_data_size, MEMREMAP_WC);
 	if (!etmdev->shared_mem_vaddr) {
 		etdev_err(etdev, "Shared memory remap failed\n");
 		ret = -ENOMEM;
 		goto out;
 	}

 	ret = edgetpu_iremap_pool_create(etdev,
 					 /* Base virtual address (kernel address space) */
 					 etmdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
 					 /* Base DMA address */
 					 etmdev->remapped_data_addr + EDGETPU_POOL_MEM_OFFSET,
 					 /* Base TPU address */
 					 etmdev->remapped_data_addr + EDGETPU_POOL_MEM_OFFSET,
 					 /* Base physical address */
 					 etmdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
 					 /* Size */
 					 etmdev->remapped_data_size - EDGETPU_POOL_MEM_OFFSET,
 					 /* Granularity */
 					 PAGE_SIZE);
 	if (ret) {
 		etdev_err(etdev, "failed to initialize remapped memory pool: %d", ret);
 		goto out_memunmap;
 	}

 	edgetpu_mobile_set_telemetry_mem(etmdev);
 	ret = edgetpu_telemetry_init(etdev, etmdev->log_mem, etmdev->trace_mem);
 	if (ret)
 		goto out_iremap_pool_destroy;

 	ret = edgetpu_kci_init(etdev->mailbox_manager, etdev->etkci);
 	if (ret)
 		goto out_telemetry_exit;

 	return 0;

 out_telemetry_exit:
 	edgetpu_telemetry_exit(etdev);
 out_iremap_pool_destroy:
 	edgetpu_iremap_pool_destroy(etdev);
 out_memunmap:
 	memunmap(etmdev->shared_mem_vaddr);

 out:
 	etmdev->remapped_data_addr = 0;
 	etmdev->remapped_data_size = 0;
 	etmdev->shared_mem_paddr = 0;
 	etmdev->shared_mem_vaddr = NULL;

 	return ret;
 }

 static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
 				       struct edgetpu_firmware_buffer *fw_buf)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;
 	int ret;

 	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
 	edgetpu_mailbox_reset(etdev->etkci->mailbox);

 	ret = mobile_firmware_update_remapped_data_region(etdev);
 	if (ret)
 		return ret;

 	edgetpu_soc_prepare_firmware(etdev);

 	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
 }

 static int mobile_firmware_restart(struct edgetpu_firmware *et_fw, bool force_reset)
 {
 	struct edgetpu_dev *etdev = et_fw->etdev;

 	/*
 	 * We are in a bad state, reset the CPU and hope the device recovers.
 	 * Ignore failures in the reset assert request and proceed to reset release.
 	 */
 	if (force_reset)
 		edgetpu_mobile_firmware_reset_cpu(etdev, true);

 	edgetpu_soc_prepare_firmware(etdev);

 	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
 }

 static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
 					struct edgetpu_firmware_buffer *fw_buf)
 {
 	int ret = 0;
 	void *image_vaddr;
 	struct edgetpu_dev *etdev = et_fw->etdev;
 	struct gcip_image_config *image_config;
 	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(et_fw);
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	phys_addr_t image_start, image_end, carveout_start, carveout_end;
 	struct mobile_image_header *hdr;

 	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
 		etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
 			  fw_buf->used_size, MOBILE_FW_HEADER_SIZE);
 		return -EINVAL;
 	}

 	hdr = (struct mobile_image_header *)fw_buf->vaddr;
 	if (hdr->common.Magic != EDGETPU_MOBILE_FW_MAGIC) {
 		etdev_err(etdev, "Invalid firmware header magic value %#08x\n", hdr->common.Magic);
 		return -EINVAL;
 	}

 	switch (hdr->common.Generation) {
 	case 1:
 		image_config = &hdr->gen1.ImageConfig;
 		break;
 	case 2:
 		image_config = &hdr->gen2.ImageConfig;
 		break;
 	default:
 		etdev_err(etdev, "Invalid header generation identifier (%d)\n",
 			  hdr->common.Generation);
 		return -EINVAL;
 	}

 	etmdev->fw_region_paddr = image_config->firmware_base;
 	etmdev->fw_region_size =
 		image_config->remapped_data_start ?
 			image_config->remapped_data_start - image_config->firmware_base :
 			EDGETPU_DEFAULT_FW_LIMIT;

 	image_vaddr = memremap(etmdev->fw_region_paddr, etmdev->fw_region_size, MEMREMAP_WC);
 	if (!image_vaddr) {
 		etdev_err(etdev, "FW region memremap failed\n");
 		return -ENOMEM;
 	}

 	memcpy(&etdev->fw_version, &image_config->firmware_versions, sizeof(etdev->fw_version));

 	/*
 	 * TODO(b/244103549) re-enable authentication of non-secure images when
 	 * GSA and TPU updates for Zuma land
 	 */
 	if (etmdev->gsa_dev && !gcip_image_config_is_ns(image_config)) {
 		ret = mobile_firmware_gsa_authenticate(etmdev, fw_buf, image_vaddr);
 	} else if (gcip_image_config_is_ns(image_config)) {
 		etdev_dbg(etdev, "Loading unauthenticated non-secure firmware\n");
 		/* Copy the firmware image to the target location, skipping the header */
 		memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
 		       fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
 	} else {
 		etdev_err(etdev,
 			  "Cannot load firmware at privilege level %d with no authentication\n",
 			  image_config->privilege_level);
 		ret = -EINVAL;
 	}

 	if (ret)
 		goto out;

 	image_start = (phys_addr_t)image_config->carveout_base;
 	image_end = (phys_addr_t)(image_config->firmware_base + image_config->firmware_size - 1);
 	carveout_start = etmdev->fw_region_paddr;
 	carveout_end = carveout_start + etmdev->fw_region_size - 1;

 	/* Image must fit within the carveout */
 	if (image_start < carveout_start || image_end > carveout_end) {
 		etdev_err(etdev, "Firmware image doesn't fit in carveout\n");
 		etdev_err(etdev, "Image config: %pap - %pap\n", &image_start, &image_end);
 		etdev_err(etdev, "Carveout: %pap - %pap\n", &carveout_start, &carveout_end);
 		ret = -ERANGE;
 		goto out;
 	}

 	ret = gcip_image_config_parse(cfg_parser, image_config);
 out:
 	memunmap(image_vaddr);
 	return ret;
 }

 /* TODO: Get the SIDs and CSRs form chip config. */
 static void program_iremap_csr(struct edgetpu_dev *etdev)
 {
 	const int ctx_id = 0, sid0 = 0x30, sid1 = 0x34;

 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY, (ctx_id << 16) | sid0);
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY + 8,
 			     (ctx_id << 16) | sid1);
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE,
 			     EDGETPU_INSTRUCTION_REMAP_BASE);
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE + 8,
 			     EDGETPU_INSTRUCTION_REMAP_BASE);

 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT,
 			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_32M);
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT + 8,
 			     EDGETPU_INSTRUCTION_REMAP_BASE + SZ_32M);

 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL, 1);
 	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL + 8, 1);
 }

 int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset)
 {
 	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
 	struct gcip_image_config *image_config = mobile_firmware_get_image_config(etdev);
 	int ret = 0;

 	if (gcip_image_config_is_ns(image_config)) {
 		int i;

 		if (!assert_reset)
 			program_iremap_csr(etdev);
 		for (i = 0; i < EDGETPU_NUM_CORES; i++)
 			edgetpu_dev_write_32_sync(etdev, EDGETPU_REG_RESET_CONTROL + i * 8,
 						  assert_reset ? 1 : 0);
 	}
 	else if (etmdev->gsa_dev)
 		ret = gsa_send_tpu_cmd(etmdev->gsa_dev,
 				       assert_reset ? GSA_TPU_SHUTDOWN : GSA_TPU_START);
 	else
 		ret = -ENODEV;

 	etdev_dbg(etdev, "%s CPU reset result = %d", assert_reset ? "assert" : "release", ret);

 	if (ret < 0)
 		return ret;

 	return 0;
 }

 static const struct edgetpu_firmware_chip_data mobile_firmware_chip_data = {
 	.default_firmware_name = EDGETPU_DEFAULT_FIRMWARE_NAME,
 	.after_create = mobile_firmware_after_create,
 	.before_destroy = mobile_firmware_before_destroy,
 	.alloc_buffer = mobile_firmware_alloc_buffer,
 	.free_buffer = mobile_firmware_free_buffer,
 	.setup_buffer = mobile_firmware_setup_buffer,
 	.prepare_run = mobile_firmware_prepare_run,
 	.restart = mobile_firmware_restart,
 };

 int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev)
 {
 	return edgetpu_firmware_create(etdev, &mobile_firmware_chip_data);
 }

 void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev)
 {
 	edgetpu_firmware_destroy(etdev);
 }

 unsigned long edgetpu_chip_firmware_iova(struct edgetpu_dev *etdev)
 {
 	/*
 	 * On mobile platforms, firmware address translation may happen in 1 or 2 stages:
 	 * 1.- Instruction remap registers.
 	 * 2.- IOMMU translation (when not running in GSA privilege).
 	 * In either case, the address seen by the TPU's CPU core will remain constant, and
 	 * equal to the macro below.
 	 */
 	return EDGETPU_INSTRUCTION_REMAP_BASE;
 }
	// SPDX-License-Identifier: GPL-2.0
	/*
	* Edge TPU firmware management for mobile chipsets.
	*
	* Copyright (C) 2021-2022 Google LLC
	*/

	#include <linux/device.h>
	#include <linux/dma-mapping.h>
	#include <linux/firmware.h>
	#include <linux/gsa/gsa_tpu.h>
	#include <linux/slab.h>
	#include <linux/string.h>
	#include <linux/types.h>

	#include <gcip/gcip-alloc-helper.h>
	#include <gcip/gcip-image-config.h>

	#include "edgetpu.h"
	#include "edgetpu-config.h"
	#include "edgetpu-firmware.h"
	#include "edgetpu-firmware-util.h"
	#include "edgetpu-internal.h"
	#include "edgetpu-iremap-pool.h"
	#include "edgetpu-kci.h"
	#include "edgetpu-mailbox.h"
	#include "edgetpu-mmu.h"
	#include "edgetpu-mobile-platform.h"
	#include "edgetpu-soc.h"
	#include "edgetpu-telemetry.h"
	#include "mobile-firmware.h"

	static int image_config_map(void *data, dma_addr_t daddr, phys_addr_t paddr, size_t size,
	unsigned int flags)
	{
	struct edgetpu_dev *etdev = data;
	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
	struct edgetpu_mobile_fw_ctx *fw_ctx;
	int ret;

	if (flags & GCIP_IMAGE_CONFIG_FLAGS_SECURE)
	return edgetpu_mmu_add_translation(etdev, daddr, paddr, size,
	IOMMU_READ \| IOMMU_WRITE, EDGETPU_CONTEXT_KCI);

	fw_ctx = kzalloc(sizeof(*fw_ctx), GFP_KERNEL);
	if (!fw_ctx)
	return -ENOMEM;

	fw_ctx->daddr = daddr;
	fw_ctx->size = size;
	fw_ctx->sgt = gcip_alloc_noncontiguous(etdev->dev, size, GFP_KERNEL);
	if (!fw_ctx->sgt) {
	kfree(fw_ctx);
	return -ENOMEM;
	}

	ret = edgetpu_mmu_map_iova_sgt(etdev, daddr, fw_ctx->sgt, DMA_BIDIRECTIONAL, 0,
	EDGETPU_CONTEXT_KCI);
	if (ret) {
	gcip_free_noncontiguous(fw_ctx->sgt);
	kfree(fw_ctx);
	return ret;
	}

	mutex_lock(&etmdev->fw_ctx_list_lock);
	list_add_tail(&fw_ctx->list, &etmdev->fw_ctx_list);
	mutex_unlock(&etmdev->fw_ctx_list_lock);

	return 0;
	}

	static void image_config_unmap(void *data, dma_addr_t daddr, size_t size, unsigned int flags)
	{
	struct edgetpu_dev *etdev = data;
	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
	struct edgetpu_mobile_fw_ctx fw_ctx = NULL, cur;

	if (flags & GCIP_IMAGE_CONFIG_FLAGS_SECURE) {
	edgetpu_mmu_remove_translation(etdev, daddr, size, EDGETPU_CONTEXT_KCI);
	return;
	}

	mutex_lock(&etmdev->fw_ctx_list_lock);
	list_for_each_entry(cur, &etmdev->fw_ctx_list, list) {
	if (cur->daddr == daddr && cur->size == size) {
	fw_ctx = cur;
	list_del(&cur->list);
	break;
	}
	}
	mutex_unlock(&etmdev->fw_ctx_list_lock);

	if (fw_ctx) {
	edgetpu_mmu_unmap_iova_sgt(etdev, daddr, fw_ctx->sgt, DMA_BIDIRECTIONAL,
	EDGETPU_CONTEXT_KCI);
	gcip_free_noncontiguous(fw_ctx->sgt);
	kfree(fw_ctx);
	} else {
	etdev_warn(etdev, "Firmware context region SG table not found.");
	}
	}

	static int mobile_firmware_after_create(struct edgetpu_firmware *et_fw)
	{
	struct gcip_image_config_parser *data;
	struct edgetpu_dev *etdev = et_fw->etdev;
	int ret;
	static const struct gcip_image_config_ops image_config_parser_ops = {
	.map = image_config_map,
	.unmap = image_config_unmap,
	};

	data = kzalloc(sizeof(*data), GFP_KERNEL);
	if (!data)
	return -ENOMEM;

	ret = gcip_image_config_parser_init(data, &image_config_parser_ops,
	get_dev_for_logging(etdev), etdev);
	if (ret) {
	etdev_err(etdev, "Image config parser init failed: %d", ret);
	kfree(data);
	return ret;
	}

	edgetpu_firmware_set_data(et_fw, data);
	return 0;
	}

	static void mobile_firmware_before_destroy(struct edgetpu_firmware *et_fw)
	{
	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(et_fw);

	gcip_image_config_clear(cfg_parser);
	edgetpu_firmware_set_data(et_fw, NULL);
	kfree(cfg_parser);
	}

	static int mobile_firmware_alloc_buffer(
	struct edgetpu_firmware *et_fw,
	struct edgetpu_firmware_buffer *fw_buf)
	{
	struct edgetpu_dev *etdev = et_fw->etdev;

	/* Allocate extra space the image header */
	size_t buffer_size = EDGETPU_MAX_FW_LIMIT + MOBILE_FW_HEADER_SIZE;

	fw_buf->vaddr = vmalloc(buffer_size);
	if (!fw_buf->vaddr) {
	etdev_err(etdev, "%s: failed to allocate buffer (%zu bytes)\n",
	__func__, buffer_size);
	return -ENOMEM;
	}
	fw_buf->dma_addr = 0;
	fw_buf->alloc_size = buffer_size;
	fw_buf->used_size_align = 16;
	return 0;
	}

	static void mobile_firmware_free_buffer(
	struct edgetpu_firmware *et_fw,
	struct edgetpu_firmware_buffer *fw_buf)
	{
	vfree(fw_buf->vaddr);
	fw_buf->vaddr = NULL;
	fw_buf->dma_addr = 0;
	fw_buf->alloc_size = 0;
	fw_buf->used_size_align = 0;
	}

	static struct gcip_image_config mobile_firmware_get_image_config(struct edgetpu_dev etdev)
	{
	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(etdev->firmware);

	return &cfg_parser->last_config;
	}

	static int mobile_firmware_gsa_authenticate(struct edgetpu_mobile_platform_dev *etmdev,
	struct edgetpu_firmware_buffer *fw_buf,
	void *image_vaddr)
	{
	struct edgetpu_dev *etdev = &etmdev->edgetpu_dev;
	void *header_vaddr;
	dma_addr_t header_dma_addr;
	int tpu_state;
	int ret = 0;

	tpu_state = gsa_send_tpu_cmd(etmdev->gsa_dev, GSA_TPU_GET_STATE);

	if (tpu_state < GSA_TPU_STATE_INACTIVE) {
	etdev_err(etdev, "GSA failed to retrieve current status: %d\n", tpu_state);
	return tpu_state;
	}

	etdev_dbg(etdev, "GSA Reports TPU state: %d\n", tpu_state);

	if (tpu_state > GSA_TPU_STATE_INACTIVE) {
	ret = gsa_unload_tpu_fw_image(etmdev->gsa_dev);
	if (ret) {
	etdev_warn(etdev, "GSA release failed: %d\n", ret);
	return -EIO;
	}
	}

	/* Copy the firmware image to the target location, skipping the header */
	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
	fw_buf->used_size - MOBILE_FW_HEADER_SIZE);

	/* Allocate coherent memory for the image header */
	header_vaddr = dma_alloc_coherent(etmdev->gsa_dev,
	MOBILE_FW_HEADER_SIZE,
	&header_dma_addr, GFP_KERNEL);
	if (!header_vaddr) {
	etdev_err(etdev,
	"Failed to allocate coherent memory for header\n");
	return -ENOMEM;
	}

	memcpy(header_vaddr, fw_buf->vaddr, MOBILE_FW_HEADER_SIZE);
	etdev_dbg(etdev, "Requesting GSA image load. meta = %llX payload = %llX", header_dma_addr,
	(u64)etmdev->fw_region_paddr);

	ret = gsa_load_tpu_fw_image(etmdev->gsa_dev, header_dma_addr,
	etmdev->fw_region_paddr);
	if (ret)
	etdev_err(etdev, "GSA authentication failed: %d\n", ret);

	dma_free_coherent(etmdev->gsa_dev, MOBILE_FW_HEADER_SIZE, header_vaddr, header_dma_addr);

	return ret;
	}

	static int mobile_firmware_update_remapped_data_region(struct edgetpu_dev *etdev)
	{
	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
	struct gcip_image_config *config = mobile_firmware_get_image_config(etdev);
	tpu_addr_t remapped_data_addr = EDGETPU_INSTRUCTION_REMAP_BASE + etmdev->fw_region_size;
	size_t remapped_data_size = config->remapped_data_size ? config->remapped_data_size :
	EDGETPU_DEFAULT_REMAPPED_DATA_SIZE;
	int ret;

	if (etmdev->remapped_data_addr == remapped_data_addr &&
	etmdev->remapped_data_size == remapped_data_size)
	return 0;

	if (remapped_data_addr < EDGETPU_INSTRUCTION_REMAP_BASE + config->firmware_size \|\|
	config->firmware_base + etmdev->fw_region_size + remapped_data_size >
	config->remapped_region_start)
	return -EINVAL;

	etdev_dbg(etdev, "Moving remapped data from %#llx to %#llx\n", etmdev->remapped_data_addr,
	remapped_data_addr);

	if (etmdev->shared_mem_vaddr) {
	/* No need to free the VII queues since allocated groups will block fw loading. */
	edgetpu_kci_release(etdev, etdev->etkci);
	edgetpu_telemetry_exit(etdev);
	edgetpu_iremap_pool_destroy(etdev);
	memunmap(etmdev->shared_mem_vaddr);
	}

	etmdev->remapped_data_addr = remapped_data_addr;
	etmdev->remapped_data_size = remapped_data_size;

	etmdev->shared_mem_paddr = etmdev->fw_region_paddr + etmdev->fw_region_size;
	etmdev->shared_mem_vaddr =
	memremap(etmdev->shared_mem_paddr, etmdev->remapped_data_size, MEMREMAP_WC);
	if (!etmdev->shared_mem_vaddr) {
	etdev_err(etdev, "Shared memory remap failed\n");
	ret = -ENOMEM;
	goto out;
	}

	ret = edgetpu_iremap_pool_create(etdev,
	/* Base virtual address (kernel address space) */
	etmdev->shared_mem_vaddr + EDGETPU_POOL_MEM_OFFSET,
	/* Base DMA address */
	etmdev->remapped_data_addr + EDGETPU_POOL_MEM_OFFSET,
	/* Base TPU address */
	etmdev->remapped_data_addr + EDGETPU_POOL_MEM_OFFSET,
	/* Base physical address */
	etmdev->shared_mem_paddr + EDGETPU_POOL_MEM_OFFSET,
	/* Size */
	etmdev->remapped_data_size - EDGETPU_POOL_MEM_OFFSET,
	/* Granularity */
	PAGE_SIZE);
	if (ret) {
	etdev_err(etdev, "failed to initialize remapped memory pool: %d", ret);
	goto out_memunmap;
	}

	edgetpu_mobile_set_telemetry_mem(etmdev);
	ret = edgetpu_telemetry_init(etdev, etmdev->log_mem, etmdev->trace_mem);
	if (ret)
	goto out_iremap_pool_destroy;

	ret = edgetpu_kci_init(etdev->mailbox_manager, etdev->etkci);
	if (ret)
	goto out_telemetry_exit;

	return 0;

	out_telemetry_exit:
	edgetpu_telemetry_exit(etdev);
	out_iremap_pool_destroy:
	edgetpu_iremap_pool_destroy(etdev);
	out_memunmap:
	memunmap(etmdev->shared_mem_vaddr);

	out:
	etmdev->remapped_data_addr = 0;
	etmdev->remapped_data_size = 0;
	etmdev->shared_mem_paddr = 0;
	etmdev->shared_mem_vaddr = NULL;

	return ret;
	}

	static int mobile_firmware_prepare_run(struct edgetpu_firmware *et_fw,
	struct edgetpu_firmware_buffer *fw_buf)
	{
	struct edgetpu_dev *etdev = et_fw->etdev;
	int ret;

	/* Reset KCI mailbox before starting f/w, don't process anything old.*/
	edgetpu_mailbox_reset(etdev->etkci->mailbox);

	ret = mobile_firmware_update_remapped_data_region(etdev);
	if (ret)
	return ret;

	edgetpu_soc_prepare_firmware(etdev);

	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
	}

	static int mobile_firmware_restart(struct edgetpu_firmware *et_fw, bool force_reset)
	{
	struct edgetpu_dev *etdev = et_fw->etdev;

	/*
	* We are in a bad state, reset the CPU and hope the device recovers.
	* Ignore failures in the reset assert request and proceed to reset release.
	*/
	if (force_reset)
	edgetpu_mobile_firmware_reset_cpu(etdev, true);

	edgetpu_soc_prepare_firmware(etdev);

	return edgetpu_mobile_firmware_reset_cpu(etdev, false);
	}

	static int mobile_firmware_setup_buffer(struct edgetpu_firmware *et_fw,
	struct edgetpu_firmware_buffer *fw_buf)
	{
	int ret = 0;
	void *image_vaddr;
	struct edgetpu_dev *etdev = et_fw->etdev;
	struct gcip_image_config *image_config;
	struct gcip_image_config_parser *cfg_parser = edgetpu_firmware_get_data(et_fw);
	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
	phys_addr_t image_start, image_end, carveout_start, carveout_end;
	struct mobile_image_header *hdr;

	if (fw_buf->used_size < MOBILE_FW_HEADER_SIZE) {
	etdev_err(etdev, "Invalid buffer size: %zu < %d\n",
	fw_buf->used_size, MOBILE_FW_HEADER_SIZE);
	return -EINVAL;
	}

	hdr = (struct mobile_image_header *)fw_buf->vaddr;
	if (hdr->common.Magic != EDGETPU_MOBILE_FW_MAGIC) {
	etdev_err(etdev, "Invalid firmware header magic value %#08x\n", hdr->common.Magic);
	return -EINVAL;
	}

	switch (hdr->common.Generation) {
	case 1:
	image_config = &hdr->gen1.ImageConfig;
	break;
	case 2:
	image_config = &hdr->gen2.ImageConfig;
	break;
	default:
	etdev_err(etdev, "Invalid header generation identifier (%d)\n",
	hdr->common.Generation);
	return -EINVAL;
	}

	etmdev->fw_region_paddr = image_config->firmware_base;
	etmdev->fw_region_size =
	image_config->remapped_data_start ?
	image_config->remapped_data_start - image_config->firmware_base :
	EDGETPU_DEFAULT_FW_LIMIT;

	image_vaddr = memremap(etmdev->fw_region_paddr, etmdev->fw_region_size, MEMREMAP_WC);
	if (!image_vaddr) {
	etdev_err(etdev, "FW region memremap failed\n");
	return -ENOMEM;
	}

	memcpy(&etdev->fw_version, &image_config->firmware_versions, sizeof(etdev->fw_version));

	/*
	* TODO(b/244103549) re-enable authentication of non-secure images when
	* GSA and TPU updates for Zuma land
	*/
	if (etmdev->gsa_dev && !gcip_image_config_is_ns(image_config)) {
	ret = mobile_firmware_gsa_authenticate(etmdev, fw_buf, image_vaddr);
	} else if (gcip_image_config_is_ns(image_config)) {
	etdev_dbg(etdev, "Loading unauthenticated non-secure firmware\n");
	/* Copy the firmware image to the target location, skipping the header */
	memcpy(image_vaddr, fw_buf->vaddr + MOBILE_FW_HEADER_SIZE,
	fw_buf->used_size - MOBILE_FW_HEADER_SIZE);
	} else {
	etdev_err(etdev,
	"Cannot load firmware at privilege level %d with no authentication\n",
	image_config->privilege_level);
	ret = -EINVAL;
	}

	if (ret)
	goto out;

	image_start = (phys_addr_t)image_config->carveout_base;
	image_end = (phys_addr_t)(image_config->firmware_base + image_config->firmware_size - 1);
	carveout_start = etmdev->fw_region_paddr;
	carveout_end = carveout_start + etmdev->fw_region_size - 1;

	/* Image must fit within the carveout */
	if (image_start < carveout_start \|\| image_end > carveout_end) {
	etdev_err(etdev, "Firmware image doesn't fit in carveout\n");
	etdev_err(etdev, "Image config: %pap - %pap\n", &image_start, &image_end);
	etdev_err(etdev, "Carveout: %pap - %pap\n", &carveout_start, &carveout_end);
	ret = -ERANGE;
	goto out;
	}

	ret = gcip_image_config_parse(cfg_parser, image_config);
	out:
	memunmap(image_vaddr);
	return ret;
	}

	/* TODO: Get the SIDs and CSRs form chip config. */
	static void program_iremap_csr(struct edgetpu_dev *etdev)
	{
	const int ctx_id = 0, sid0 = 0x30, sid1 = 0x34;

	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY, (ctx_id << 16) \| sid0);
	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_SECURITY + 8,
	(ctx_id << 16) \| sid1);
	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE,
	EDGETPU_INSTRUCTION_REMAP_BASE);
	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_NEW_BASE + 8,
	EDGETPU_INSTRUCTION_REMAP_BASE);

	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT,
	EDGETPU_INSTRUCTION_REMAP_BASE + SZ_32M);
	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_LIMIT + 8,
	EDGETPU_INSTRUCTION_REMAP_BASE + SZ_32M);

	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL, 1);
	edgetpu_dev_write_32(etdev, EDGETPU_REG_INSTRUCTION_REMAP_CONTROL + 8, 1);
	}

	int edgetpu_mobile_firmware_reset_cpu(struct edgetpu_dev *etdev, bool assert_reset)
	{
	struct edgetpu_mobile_platform_dev *etmdev = to_mobile_dev(etdev);
	struct gcip_image_config *image_config = mobile_firmware_get_image_config(etdev);
	int ret = 0;

	if (gcip_image_config_is_ns(image_config)) {
	int i;

	if (!assert_reset)
	program_iremap_csr(etdev);
	for (i = 0; i < EDGETPU_NUM_CORES; i++)
	edgetpu_dev_write_32_sync(etdev, EDGETPU_REG_RESET_CONTROL + i * 8,
	assert_reset ? 1 : 0);
	}
	else if (etmdev->gsa_dev)
	ret = gsa_send_tpu_cmd(etmdev->gsa_dev,
	assert_reset ? GSA_TPU_SHUTDOWN : GSA_TPU_START);
	else
	ret = -ENODEV;

	etdev_dbg(etdev, "%s CPU reset result = %d", assert_reset ? "assert" : "release", ret);

	if (ret < 0)
	return ret;

	return 0;
	}

	static const struct edgetpu_firmware_chip_data mobile_firmware_chip_data = {
	.default_firmware_name = EDGETPU_DEFAULT_FIRMWARE_NAME,
	.after_create = mobile_firmware_after_create,
	.before_destroy = mobile_firmware_before_destroy,
	.alloc_buffer = mobile_firmware_alloc_buffer,
	.free_buffer = mobile_firmware_free_buffer,
	.setup_buffer = mobile_firmware_setup_buffer,
	.prepare_run = mobile_firmware_prepare_run,
	.restart = mobile_firmware_restart,
	};

	int edgetpu_mobile_firmware_create(struct edgetpu_dev *etdev)
	{
	return edgetpu_firmware_create(etdev, &mobile_firmware_chip_data);
	}

	void edgetpu_mobile_firmware_destroy(struct edgetpu_dev *etdev)
	{
	edgetpu_firmware_destroy(etdev);
	}

	unsigned long edgetpu_chip_firmware_iova(struct edgetpu_dev *etdev)
	{
	/*
	* On mobile platforms, firmware address translation may happen in 1 or 2 stages:
	* 1.- Instruction remap registers.
	* 2.- IOMMU translation (when not running in GSA privilege).
	* In either case, the address seen by the TPU's CPU core will remain constant, and
	* equal to the macro below.
	*/
	return EDGETPU_INSTRUCTION_REMAP_BASE;
	}