// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2026, Advanced Micro Devices, Inc.
 */

#include <drm/drm_mm.h>
#include <drm/drm_prime.h>

#include "amdxdna_cbuf.h"
#include "amdxdna_pci_drv.h"

/*
 * Carveout memory is a chunk of memory which is physically contiguous and
 * is reserved during early boot time. There is only one chunk of such memory
 * per device. Once available, all BOs accessible from device should be
 * allocated from this memory. This is a platform debug/bringup feature.
 */
struct amdxdna_carveout {
	u64		addr;
	u64		size;
	struct drm_mm	mm;
	struct mutex	lock; /* protect mm */
};

bool amdxdna_use_carveout(struct amdxdna_dev *xdna)
{
	return !!xdna->carveout;
}

void amdxdna_get_carveout_conf(struct amdxdna_dev *xdna, u64 *addr, u64 *size)
{
	if (amdxdna_use_carveout(xdna)) {
		*addr = xdna->carveout->addr;
		*size = xdna->carveout->size;
	} else {
		*addr = 0;
		*size = 0;
	}
}

int amdxdna_carveout_init(struct amdxdna_dev *xdna, u64 carveout_addr, u64 carveout_size)
{
	struct amdxdna_carveout *carveout;

	/* Only allow carveout memory to be set up once. */
	if (amdxdna_use_carveout(xdna)) {
		XDNA_ERR(xdna, "Carveout memory has already been set up.");
		return -EBUSY;
	}

	carveout = kzalloc_obj(*carveout);
	if (!carveout)
		return -ENOMEM;

	carveout->addr = carveout_addr;
	carveout->size = carveout_size;
	mutex_init(&carveout->lock);
	drm_mm_init(&carveout->mm, carveout->addr, carveout->size);

	xdna->carveout = carveout;
	XDNA_INFO(xdna, "Use carveout mem: 0x%llx@0x%llx\n", carveout->size, carveout->addr);
	return 0;
}

void amdxdna_carveout_fini(struct amdxdna_dev *xdna)
{
	struct amdxdna_carveout *carveout = xdna->carveout;

	if (!amdxdna_use_carveout(xdna))
		return;

	XDNA_INFO(xdna, "Cleanup carveout mem: 0x%llx@0x%llx\n", carveout->size, carveout->addr);
	drm_mm_takedown(&carveout->mm);
	mutex_destroy(&carveout->lock);
	kfree(carveout);
	xdna->carveout = NULL;
}

struct amdxdna_cbuf_priv {
	struct amdxdna_dev *xdna;
	struct drm_mm_node node;
};

static struct sg_table *amdxdna_cbuf_map(struct dma_buf_attachment *attach,
					 enum dma_data_direction direction)
{
	struct amdxdna_cbuf_priv *cbuf = attach->dmabuf->priv;
	struct device *dev = attach->dev;
	struct scatterlist *sgl, *sg;
	int ret, n_entries, i;
	struct sg_table *sgt;
	dma_addr_t dma_addr;
	size_t dma_size;
	size_t max_seg;

	sgt = kzalloc_obj(*sgt);
	if (!sgt)
		return ERR_PTR(-ENOMEM);

	max_seg = min_t(size_t, UINT_MAX, dma_max_mapping_size(dev));
	n_entries = (cbuf->node.size + max_seg - 1) / max_seg;
	sgl = kzalloc_objs(*sg, n_entries);
	if (!sgl) {
		ret = -ENOMEM;
		goto free_sgt;
	}
	sg_init_table(sgl, n_entries);
	sgt->orig_nents = n_entries;
	sgt->nents = n_entries;
	sgt->sgl = sgl;

	dma_size = cbuf->node.size;
	dma_addr = dma_map_resource(dev, cbuf->node.start, dma_size,
				    direction, DMA_ATTR_SKIP_CPU_SYNC);
	ret = dma_mapping_error(dev, dma_addr);
	if (ret) {
		pr_err("Failed to dma_map_resource carveout dma buf, ret %d\n", ret);
		goto free_sgl;
	}

	for_each_sgtable_dma_sg(sgt, sg, i) {
		size_t len = min_t(size_t, max_seg, dma_size);

		sg_dma_address(sg) = dma_addr;
		sg_dma_len(sg) = len;
		dma_addr += len;
		dma_size -= len;
	}

	return sgt;

free_sgl:
	kfree(sgl);
free_sgt:
	kfree(sgt);
	return ERR_PTR(ret);
}

static void amdxdna_cbuf_unmap(struct dma_buf_attachment *attach,
			       struct sg_table *sgt,
			       enum dma_data_direction direction)
{
	dma_unmap_resource(attach->dev, sg_dma_address(sgt->sgl),
			   drm_prime_get_contiguous_size(sgt), direction,
			   DMA_ATTR_SKIP_CPU_SYNC);
	sg_free_table(sgt);
	kfree(sgt);
}

static void amdxdna_cbuf_release(struct dma_buf *dbuf)
{
	struct amdxdna_cbuf_priv *cbuf = dbuf->priv;
	struct amdxdna_carveout *carveout;

	carveout = cbuf->xdna->carveout;
	mutex_lock(&carveout->lock);
	drm_mm_remove_node(&cbuf->node);
	mutex_unlock(&carveout->lock);

	kfree(cbuf);
}

static vm_fault_t amdxdna_cbuf_vm_fault(struct vm_fault *vmf)
{
	struct vm_area_struct *vma = vmf->vma;
	struct amdxdna_cbuf_priv *cbuf;
	unsigned long pfn;
	pgoff_t pgoff;

	cbuf = vma->vm_private_data;
	pgoff = (vmf->address - vma->vm_start) >> PAGE_SHIFT;
	pfn = (cbuf->node.start >> PAGE_SHIFT) + pgoff;

	return vmf_insert_pfn(vma, vmf->address, pfn);
}

static const struct vm_operations_struct amdxdna_cbuf_vm_ops = {
	.fault = amdxdna_cbuf_vm_fault,
};

static int amdxdna_cbuf_mmap(struct dma_buf *dbuf, struct vm_area_struct *vma)
{
	struct amdxdna_cbuf_priv *cbuf = dbuf->priv;

	vma->vm_ops = &amdxdna_cbuf_vm_ops;
	vma->vm_private_data = cbuf;
	vm_flags_set(vma, VM_PFNMAP | VM_DONTEXPAND | VM_DONTDUMP);

	return 0;
}

static int amdxdna_cbuf_vmap(struct dma_buf *dbuf, struct iosys_map *map)
{
	struct amdxdna_cbuf_priv *cbuf = dbuf->priv;
	void *kva;

	kva = memremap(cbuf->node.start, cbuf->node.size, MEMREMAP_WB);
	if (!kva) {
		pr_err("Failed to vmap carveout dma buf\n");
		return -ENOMEM;
	}

	iosys_map_set_vaddr(map, kva);
	return 0;
}

static void amdxdna_cbuf_vunmap(struct dma_buf *dbuf, struct iosys_map *map)
{
	memunmap(map->vaddr);
}

static const struct dma_buf_ops amdxdna_cbuf_dmabuf_ops = {
	.map_dma_buf = amdxdna_cbuf_map,
	.unmap_dma_buf = amdxdna_cbuf_unmap,
	.release = amdxdna_cbuf_release,
	.mmap = amdxdna_cbuf_mmap,
	.vmap = amdxdna_cbuf_vmap,
	.vunmap = amdxdna_cbuf_vunmap,
};

static int amdxdna_cbuf_clear(struct dma_buf *dbuf)
{
	struct iosys_map vmap = IOSYS_MAP_INIT_VADDR(NULL);

	dma_buf_vmap(dbuf, &vmap);
	if (!vmap.vaddr)
		return -EFAULT;

	memset(vmap.vaddr, 0, dbuf->size);
	dma_buf_vunmap(dbuf, &vmap);

	return 0;
}

struct dma_buf *amdxdna_get_cbuf(struct drm_device *dev, size_t size, u64 alignment)
{
	struct amdxdna_dev *xdna = to_xdna_dev(dev);
	DEFINE_DMA_BUF_EXPORT_INFO(exp_info);
	struct amdxdna_carveout *carveout;
	struct amdxdna_cbuf_priv *cbuf;
	struct dma_buf *dbuf;
	int ret;

	cbuf = kzalloc_obj(*cbuf);
	if (!cbuf)
		return ERR_PTR(-ENOMEM);
	cbuf->xdna = xdna;

	carveout = xdna->carveout;
	mutex_lock(&carveout->lock);
	ret = drm_mm_insert_node_generic(&carveout->mm, &cbuf->node, size,
					 alignment, 0, DRM_MM_INSERT_BEST);
	mutex_unlock(&carveout->lock);
	if (ret)
		goto free_cbuf;

	exp_info.size = size;
	exp_info.ops = &amdxdna_cbuf_dmabuf_ops;
	exp_info.priv = cbuf;
	exp_info.flags = O_RDWR;
	dbuf = dma_buf_export(&exp_info);
	if (IS_ERR(dbuf)) {
		ret = PTR_ERR(dbuf);
		goto remove_node;
	}

	ret = amdxdna_cbuf_clear(dbuf);
	if (ret) {
		dma_buf_put(dbuf);
		goto out;
	}
	return dbuf;

remove_node:
	drm_mm_remove_node(&cbuf->node);
free_cbuf:
	kfree(cbuf);
out:
	return ERR_PTR(ret);
}