// SPDX-License-Identifier: GPL-2.0 or MIT
/* Copyright 2023 Collabora ltd. */

#include <linux/iosys-map.h>
#include <linux/rwsem.h>

#include <drm/panthor_drm.h>

#include "panthor_device.h"
#include "panthor_gem.h"
#include "panthor_heap.h"
#include "panthor_mmu.h"
#include "panthor_regs.h"

/*
 * The GPU heap context is an opaque structure used by the GPU to track the
 * heap allocations. The driver should only touch it to initialize it (zero all
 * fields). Because the CPU and GPU can both access this structure it is
 * required to be GPU cache line aligned.
 */
#define HEAP_CONTEXT_SIZE	32

/**
 * struct panthor_heap_chunk_header - Heap chunk header
 */
struct panthor_heap_chunk_header {
	/**
	 * @next: Next heap chunk in the list.
	 *
	 * This is a GPU VA.
	 */
	u64 next;

	/** @unknown: MBZ. */
	u32 unknown[14];
};

/**
 * struct panthor_heap_chunk - Structure used to keep track of allocated heap chunks.
 */
struct panthor_heap_chunk {
	/** @node: Used to insert the heap chunk in panthor_heap::chunks. */
	struct list_head node;

	/** @bo: Buffer object backing the heap chunk. */
	struct panthor_kernel_bo *bo;
};

/**
 * struct panthor_heap - Structure used to manage tiler heap contexts.
 */
struct panthor_heap {
	/** @chunks: List containing all heap chunks allocated so far. */
	struct list_head chunks;

	/** @lock: Lock protecting insertion in the chunks list. */
	struct mutex lock;

	/** @chunk_size: Size of each chunk. */
	u32 chunk_size;

	/** @max_chunks: Maximum number of chunks. */
	u32 max_chunks;

	/**
	 * @target_in_flight: Number of in-flight render passes after which
	 * we'd let the FW wait for fragment job to finish instead of allocating new chunks.
	 */
	u32 target_in_flight;

	/** @chunk_count: Number of heap chunks currently allocated. */
	u32 chunk_count;
};

#define MAX_HEAPS_PER_POOL    128

/**
 * struct panthor_heap_pool - Pool of heap contexts
 *
 * The pool is attached to a panthor_file and can't be shared across processes.
 */
struct panthor_heap_pool {
	/** @refcount: Reference count. */
	struct kref refcount;

	/** @ptdev: Device. */
	struct panthor_device *ptdev;

	/** @vm: VM this pool is bound to. */
	struct panthor_vm *vm;

	/** @lock: Lock protecting access to @xa. */
	struct rw_semaphore lock;

	/** @xa: Array storing panthor_heap objects. */
	struct xarray xa;

	/** @gpu_contexts: Buffer object containing the GPU heap contexts. */
	struct panthor_kernel_bo *gpu_contexts;

	/** @size: Size of all chunks across all heaps in the pool. */
	atomic_t size;
};

static int panthor_heap_ctx_stride(struct panthor_device *ptdev)
{
	u32 l2_features = ptdev->gpu_info.l2_features;
	u32 gpu_cache_line_size = GPU_L2_FEATURES_LINE_SIZE(l2_features);

	return ALIGN(HEAP_CONTEXT_SIZE, gpu_cache_line_size);
}

static int panthor_get_heap_ctx_offset(struct panthor_heap_pool *pool, int id)
{
	return panthor_heap_ctx_stride(pool->ptdev) * id;
}

static void *panthor_get_heap_ctx(struct panthor_heap_pool *pool, int id)
{
	return pool->gpu_contexts->kmap +
	       panthor_get_heap_ctx_offset(pool, id);
}

static void panthor_free_heap_chunk(struct panthor_heap_pool *pool,
				    struct panthor_heap *heap,
				    struct panthor_heap_chunk *chunk)
{
	mutex_lock(&heap->lock);
	list_del(&chunk->node);
	heap->chunk_count--;
	mutex_unlock(&heap->lock);

	atomic_sub(heap->chunk_size, &pool->size);

	panthor_kernel_bo_destroy(chunk->bo);
	kfree(chunk);
}

static int panthor_alloc_heap_chunk(struct panthor_heap_pool *pool,
				    struct panthor_heap *heap,
				    bool initial_chunk)
{
	struct panthor_heap_chunk *chunk;
	struct panthor_heap_chunk_header *hdr;
	int ret;

	chunk = kmalloc(sizeof(*chunk), GFP_KERNEL);
	if (!chunk)
		return -ENOMEM;

	chunk->bo = panthor_kernel_bo_create(pool->ptdev, pool->vm, heap->chunk_size,
					     DRM_PANTHOR_BO_NO_MMAP,
					     DRM_PANTHOR_VM_BIND_OP_MAP_NOEXEC,
					     PANTHOR_VM_KERNEL_AUTO_VA,
					     "Tiler heap chunk");
	if (IS_ERR(chunk->bo)) {
		ret = PTR_ERR(chunk->bo);
		goto err_free_chunk;
	}

	ret = panthor_kernel_bo_vmap(chunk->bo);
	if (ret)
		goto err_destroy_bo;

	hdr = chunk->bo->kmap;
	memset(hdr, 0, sizeof(*hdr));

	if (initial_chunk && !list_empty(&heap->chunks)) {
		struct panthor_heap_chunk *prev_chunk;
		u64 prev_gpuva;

		prev_chunk = list_first_entry(&heap->chunks,
					      struct panthor_heap_chunk,
					      node);

		prev_gpuva = panthor_kernel_bo_gpuva(prev_chunk->bo);
		hdr->next = (prev_gpuva & GENMASK_ULL(63, 12)) |
			    (heap->chunk_size >> 12);
	}

	panthor_kernel_bo_vunmap(chunk->bo);

	mutex_lock(&heap->lock);
	list_add(&chunk->node, &heap->chunks);
	heap->chunk_count++;
	mutex_unlock(&heap->lock);

	atomic_add(heap->chunk_size, &pool->size);

	return 0;

err_destroy_bo:
	panthor_kernel_bo_destroy(chunk->bo);

err_free_chunk:
	kfree(chunk);

	return ret;
}

static void panthor_free_heap_chunks(struct panthor_heap_pool *pool,
				     struct panthor_heap *heap)
{
	struct panthor_heap_chunk *chunk, *tmp;

	list_for_each_entry_safe(chunk, tmp, &heap->chunks, node)
		panthor_free_heap_chunk(pool, heap, chunk);
}

static int panthor_alloc_heap_chunks(struct panthor_heap_pool *pool,
				     struct panthor_heap *heap,
				     u32 chunk_count)
{
	int ret;
	u32 i;

	for (i = 0; i < chunk_count; i++) {
		ret = panthor_alloc_heap_chunk(pool, heap, true);
		if (ret)
			return ret;
	}

	return 0;
}

static int
panthor_heap_destroy_locked(struct panthor_heap_pool *pool, u32 handle)
{
	struct panthor_heap *heap;

	heap = xa_erase(&pool->xa, handle);
	if (!heap)
		return -EINVAL;

	panthor_free_heap_chunks(pool, heap);
	mutex_destroy(&heap->lock);
	kfree(heap);
	return 0;
}

/**
 * panthor_heap_destroy() - Destroy a heap context
 * @pool: Pool this context belongs to.
 * @handle: Handle returned by panthor_heap_create().
 */
int panthor_heap_destroy(struct panthor_heap_pool *pool, u32 handle)
{
	int ret;

	down_write(&pool->lock);
	ret = panthor_heap_destroy_locked(pool, handle);
	up_write(&pool->lock);

	return ret;
}

/**
 * panthor_heap_create() - Create a heap context
 * @pool: Pool to instantiate the heap context from.
 * @initial_chunk_count: Number of chunk allocated at initialization time.
 * Must be at least 1.
 * @chunk_size: The size of each chunk. Must be page-aligned and lie in the
 * [128k:8M] range.
 * @max_chunks: Maximum number of chunks that can be allocated.
 * @target_in_flight: Maximum number of in-flight render passes.
 * @heap_ctx_gpu_va: Pointer holding the GPU address of the allocated heap
 * context.
 * @first_chunk_gpu_va: Pointer holding the GPU address of the first chunk
 * assigned to the heap context.
 *
 * Return: a positive handle on success, a negative error otherwise.
 */
int panthor_heap_create(struct panthor_heap_pool *pool,
			u32 initial_chunk_count,
			u32 chunk_size,
			u32 max_chunks,
			u32 target_in_flight,
			u64 *heap_ctx_gpu_va,
			u64 *first_chunk_gpu_va)
{
	struct panthor_heap *heap;
	struct panthor_heap_chunk *first_chunk;
	struct panthor_vm *vm;
	int ret = 0;
	u32 id;

	if (initial_chunk_count == 0)
		return -EINVAL;

	if (initial_chunk_count > max_chunks)
		return -EINVAL;

	if (!IS_ALIGNED(chunk_size, PAGE_SIZE) ||
	    chunk_size < SZ_128K || chunk_size > SZ_8M)
		return -EINVAL;

	down_read(&pool->lock);
	vm = panthor_vm_get(pool->vm);
	up_read(&pool->lock);

	/* The pool has been destroyed, we can't create a new heap. */
	if (!vm)
		return -EINVAL;

	heap = kzalloc(sizeof(*heap), GFP_KERNEL);
	if (!heap) {
		ret = -ENOMEM;
		goto err_put_vm;
	}

	mutex_init(&heap->lock);
	INIT_LIST_HEAD(&heap->chunks);
	heap->chunk_size = chunk_size;
	heap->max_chunks = max_chunks;
	heap->target_in_flight = target_in_flight;

	ret = panthor_alloc_heap_chunks(pool, heap, initial_chunk_count);
	if (ret)
		goto err_free_heap;

	first_chunk = list_first_entry(&heap->chunks,
				       struct panthor_heap_chunk,
				       node);
	*first_chunk_gpu_va = panthor_kernel_bo_gpuva(first_chunk->bo);

	down_write(&pool->lock);
	/* The pool has been destroyed, we can't create a new heap. */
	if (!pool->vm) {
		ret = -EINVAL;
	} else {
		ret = xa_alloc(&pool->xa, &id, heap,
			       XA_LIMIT(0, MAX_HEAPS_PER_POOL - 1), GFP_KERNEL);
		if (!ret) {
			void *gpu_ctx = panthor_get_heap_ctx(pool, id);

			memset(gpu_ctx, 0, panthor_heap_ctx_stride(pool->ptdev));
			*heap_ctx_gpu_va = panthor_kernel_bo_gpuva(pool->gpu_contexts) +
					   panthor_get_heap_ctx_offset(pool, id);
		}
	}
	up_write(&pool->lock);

	if (ret)
		goto err_free_heap;

	panthor_vm_put(vm);
	return id;

err_free_heap:
	panthor_free_heap_chunks(pool, heap);
	mutex_destroy(&heap->lock);
	kfree(heap);

err_put_vm:
	panthor_vm_put(vm);
	return ret;
}

/**
 * panthor_heap_return_chunk() - Return an unused heap chunk
 * @pool: The pool this heap belongs to.
 * @heap_gpu_va: The GPU address of the heap context.
 * @chunk_gpu_va: The chunk VA to return.
 *
 * This function is used when a chunk allocated with panthor_heap_grow()
 * couldn't be linked to the heap context through the FW interface because
 * the group requesting the allocation was scheduled out in the meantime.
 */
int panthor_heap_return_chunk(struct panthor_heap_pool *pool,
			      u64 heap_gpu_va,
			      u64 chunk_gpu_va)
{
	u64 offset = heap_gpu_va - panthor_kernel_bo_gpuva(pool->gpu_contexts);
	u32 heap_id = (u32)offset / panthor_heap_ctx_stride(pool->ptdev);
	struct panthor_heap_chunk *chunk, *tmp, *removed = NULL;
	struct panthor_heap *heap;
	int ret;

	if (offset > U32_MAX || heap_id >= MAX_HEAPS_PER_POOL)
		return -EINVAL;

	down_read(&pool->lock);
	heap = xa_load(&pool->xa, heap_id);
	if (!heap) {
		ret = -EINVAL;
		goto out_unlock;
	}

	chunk_gpu_va &= GENMASK_ULL(63, 12);

	mutex_lock(&heap->lock);
	list_for_each_entry_safe(chunk, tmp, &heap->chunks, node) {
		if (panthor_kernel_bo_gpuva(chunk->bo) == chunk_gpu_va) {
			removed = chunk;
			list_del(&chunk->node);
			heap->chunk_count--;
			atomic_sub(heap->chunk_size, &pool->size);
			break;
		}
	}
	mutex_unlock(&heap->lock);

	if (removed) {
		panthor_kernel_bo_destroy(chunk->bo);
		kfree(chunk);
		ret = 0;
	} else {
		ret = -EINVAL;
	}

out_unlock:
	up_read(&pool->lock);
	return ret;
}

/**
 * panthor_heap_grow() - Make a heap context grow.
 * @pool: The pool this heap belongs to.
 * @heap_gpu_va: The GPU address of the heap context.
 * @renderpasses_in_flight: Number of render passes currently in-flight.
 * @pending_frag_count: Number of fragment jobs waiting for execution/completion.
 * @new_chunk_gpu_va: Pointer used to return the chunk VA.
 *
 * Return:
 * - 0 if a new heap was allocated
 * - -ENOMEM if the tiler context reached the maximum number of chunks
 *   or if too many render passes are in-flight
 *   or if the allocation failed
 * - -EINVAL if any of the arguments passed to panthor_heap_grow() is invalid
 */
int panthor_heap_grow(struct panthor_heap_pool *pool,
		      u64 heap_gpu_va,
		      u32 renderpasses_in_flight,
		      u32 pending_frag_count,
		      u64 *new_chunk_gpu_va)
{
	u64 offset = heap_gpu_va - panthor_kernel_bo_gpuva(pool->gpu_contexts);
	u32 heap_id = (u32)offset / panthor_heap_ctx_stride(pool->ptdev);
	struct panthor_heap_chunk *chunk;
	struct panthor_heap *heap;
	int ret;

	if (offset > U32_MAX || heap_id >= MAX_HEAPS_PER_POOL)
		return -EINVAL;

	down_read(&pool->lock);
	heap = xa_load(&pool->xa, heap_id);
	if (!heap) {
		ret = -EINVAL;
		goto out_unlock;
	}

	/* If we reached the target in-flight render passes, or if we
	 * reached the maximum number of chunks, let the FW figure another way to
	 * find some memory (wait for render passes to finish, or call the exception
	 * handler provided by the userspace driver, if any).
	 */
	if (renderpasses_in_flight > heap->target_in_flight ||
	    heap->chunk_count >= heap->max_chunks) {
		ret = -ENOMEM;
		goto out_unlock;
	}

	/* FIXME: panthor_alloc_heap_chunk() triggers a kernel BO creation,
	 * which goes through the blocking allocation path. Ultimately, we
	 * want a non-blocking allocation, so we can immediately report to the
	 * FW when the system is running out of memory. In that case, the FW
	 * can call a user-provided exception handler, which might try to free
	 * some tiler memory by issuing an intermediate fragment job. If the
	 * exception handler can't do anything, it will flag the queue as
	 * faulty so the job that triggered this tiler chunk allocation and all
	 * further jobs in this queue fail immediately instead of having to
	 * wait for the job timeout.
	 */
	ret = panthor_alloc_heap_chunk(pool, heap, false);
	if (ret)
		goto out_unlock;

	chunk = list_first_entry(&heap->chunks,
				 struct panthor_heap_chunk,
				 node);
	*new_chunk_gpu_va = (panthor_kernel_bo_gpuva(chunk->bo) & GENMASK_ULL(63, 12)) |
			    (heap->chunk_size >> 12);
	ret = 0;

out_unlock:
	up_read(&pool->lock);
	return ret;
}

static void panthor_heap_pool_release(struct kref *refcount)
{
	struct panthor_heap_pool *pool =
		container_of(refcount, struct panthor_heap_pool, refcount);

	xa_destroy(&pool->xa);
	kfree(pool);
}

/**
 * panthor_heap_pool_put() - Release a heap pool reference
 * @pool: Pool to release the reference on. Can be NULL.
 */
void panthor_heap_pool_put(struct panthor_heap_pool *pool)
{
	if (pool)
		kref_put(&pool->refcount, panthor_heap_pool_release);
}

/**
 * panthor_heap_pool_get() - Get a heap pool reference
 * @pool: Pool to get the reference on. Can be NULL.
 *
 * Return: @pool.
 */
struct panthor_heap_pool *
panthor_heap_pool_get(struct panthor_heap_pool *pool)
{
	if (pool)
		kref_get(&pool->refcount);

	return pool;
}

/**
 * panthor_heap_pool_create() - Create a heap pool
 * @ptdev: Device.
 * @vm: The VM this heap pool will be attached to.
 *
 * Heap pools might contain up to 128 heap contexts, and are per-VM.
 *
 * Return: A valid pointer on success, a negative error code otherwise.
 */
struct panthor_heap_pool *
panthor_heap_pool_create(struct panthor_device *ptdev, struct panthor_vm *vm)
{
	size_t bosize = ALIGN(MAX_HEAPS_PER_POOL *
			      panthor_heap_ctx_stride(ptdev),
			      4096);
	struct panthor_heap_pool *pool;
	int ret = 0;

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

	/* We want a weak ref here: the heap pool belongs to the VM, so we're
	 * sure that, as long as the heap pool exists, the VM exists too.
	 */
	pool->vm = vm;
	pool->ptdev = ptdev;
	init_rwsem(&pool->lock);
	xa_init_flags(&pool->xa, XA_FLAGS_ALLOC);
	kref_init(&pool->refcount);

	pool->gpu_contexts = panthor_kernel_bo_create(ptdev, vm, bosize,
						      DRM_PANTHOR_BO_NO_MMAP,
						      DRM_PANTHOR_VM_BIND_OP_MAP_NOEXEC,
						      PANTHOR_VM_KERNEL_AUTO_VA,
						      "Heap pool");
	if (IS_ERR(pool->gpu_contexts)) {
		ret = PTR_ERR(pool->gpu_contexts);
		goto err_destroy_pool;
	}

	ret = panthor_kernel_bo_vmap(pool->gpu_contexts);
	if (ret)
		goto err_destroy_pool;

	atomic_add(pool->gpu_contexts->obj->size, &pool->size);

	return pool;

err_destroy_pool:
	panthor_heap_pool_destroy(pool);
	return ERR_PTR(ret);
}

/**
 * panthor_heap_pool_destroy() - Destroy a heap pool.
 * @pool: Pool to destroy.
 *
 * This function destroys all heap contexts and their resources. Thus
 * preventing any use of the heap context or the chunk attached to them
 * after that point.
 *
 * If the GPU still has access to some heap contexts, a fault should be
 * triggered, which should flag the command stream groups using these
 * context as faulty.
 *
 * The heap pool object is only released when all references to this pool
 * are released.
 */
void panthor_heap_pool_destroy(struct panthor_heap_pool *pool)
{
	struct panthor_heap *heap;
	unsigned long i;

	if (!pool)
		return;

	down_write(&pool->lock);
	xa_for_each(&pool->xa, i, heap)
		drm_WARN_ON(&pool->ptdev->base, panthor_heap_destroy_locked(pool, i));

	if (!IS_ERR_OR_NULL(pool->gpu_contexts)) {
		atomic_sub(pool->gpu_contexts->obj->size, &pool->size);
		panthor_kernel_bo_destroy(pool->gpu_contexts);
	}

	/* Reflects the fact the pool has been destroyed. */
	pool->vm = NULL;
	up_write(&pool->lock);

	panthor_heap_pool_put(pool);
}

/**
 * panthor_heap_pool_size() - Get a heap pool's total size
 * @pool: Pool whose total chunks size to return
 *
 * Returns the aggregated size of all chunks for all heaps in the pool
 *
 */
size_t panthor_heap_pool_size(struct panthor_heap_pool *pool)
{
	if (!pool)
		return 0;

	return atomic_read(&pool->size);
}