/* SPDX-License-Identifier: MIT
 *
 * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 */
#include "vmm.h"

#include <subdev/fb.h>

#include <nvhw/drf.h>
#include <nvhw/ref/gh100/dev_mmu.h>

static inline void
gh100_vmm_pgt_pte(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes,
		  struct nvkm_vmm_map *map, u64 addr)
{
	u64 data = addr | map->type;

	while (ptes--) {
		VMM_WO064(pt, vmm, ptei++ * NV_MMU_VER3_PTE__SIZE, data);
		data += map->next;
	}
}

static void
gh100_vmm_pgt_sgl(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes,
		  struct nvkm_vmm_map *map)
{
	VMM_MAP_ITER_SGL(vmm, pt, ptei, ptes, map, gh100_vmm_pgt_pte);
}

static void
gh100_vmm_pgt_dma(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes,
		  struct nvkm_vmm_map *map)
{
	if (map->page->shift == PAGE_SHIFT) {
		VMM_SPAM(vmm, "DMAA %08x %08x PTE(s)", ptei, ptes);

		nvkm_kmap(pt->memory);
		while (ptes--) {
			const u64 data = *map->dma++ | map->type;

			VMM_WO064(pt, vmm, ptei++ * NV_MMU_VER3_PTE__SIZE, data);
		}
		nvkm_done(pt->memory);
		return;
	}

	VMM_MAP_ITER_DMA(vmm, pt, ptei, ptes, map, gh100_vmm_pgt_pte);
}

static void
gh100_vmm_pgt_mem(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes,
		  struct nvkm_vmm_map *map)
{
	VMM_MAP_ITER_MEM(vmm, pt, ptei, ptes, map, gh100_vmm_pgt_pte);
}

static void
gh100_vmm_pgt_sparse(struct nvkm_vmm *vmm,
		     struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes)
{
	const u64 data = NVDEF(NV_MMU, VER3_PTE, PCF, SPARSE);

	VMM_FO064(pt, vmm, ptei * NV_MMU_VER3_PTE__SIZE, data, ptes);
}

static const struct nvkm_vmm_desc_func
gh100_vmm_desc_spt = {
	.unmap = gf100_vmm_pgt_unmap,
	.sparse = gh100_vmm_pgt_sparse,
	.mem = gh100_vmm_pgt_mem,
	.dma = gh100_vmm_pgt_dma,
	.sgl = gh100_vmm_pgt_sgl,
};

static void
gh100_vmm_lpt_invalid(struct nvkm_vmm *vmm,
		      struct nvkm_mmu_pt *pt, u32 ptei, u32 ptes)
{
	const u64 data = NVDEF(NV_MMU, VER3_PTE, PCF, NO_VALID_4KB_PAGE);

	VMM_FO064(pt, vmm, ptei * NV_MMU_VER3_PTE__SIZE, data, ptes);
}

static const struct nvkm_vmm_desc_func
gh100_vmm_desc_lpt = {
	.invalid = gh100_vmm_lpt_invalid,
	.unmap = gf100_vmm_pgt_unmap,
	.sparse = gh100_vmm_pgt_sparse,
	.mem = gh100_vmm_pgt_mem,
};

static inline void
gh100_vmm_pd0_pte(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt,
		  u32 ptei, u32 ptes, struct nvkm_vmm_map *map, u64 addr)
{
	u64 data = addr | map->type;

	while (ptes--) {
		VMM_WO128(pt, vmm, ptei++ * NV_MMU_VER3_DUAL_PDE__SIZE, data, 0ULL);
		data += map->next;
	}
}

static void
gh100_vmm_pd0_mem(struct nvkm_vmm *vmm, struct nvkm_mmu_pt *pt,
		  u32 ptei, u32 ptes, struct nvkm_vmm_map *map)
{
	VMM_MAP_ITER_MEM(vmm, pt, ptei, ptes, map, gh100_vmm_pd0_pte);
}

static inline bool
gh100_vmm_pde(struct nvkm_mmu_pt *pt, u64 *data)
{
	switch (nvkm_memory_target(pt->memory)) {
	case NVKM_MEM_TARGET_VRAM:
		*data |= NVDEF(NV_MMU, VER3_PDE, APERTURE, VIDEO_MEMORY);
		*data |= NVDEF(NV_MMU, VER3_PDE, PCF, VALID_CACHED_ATS_NOT_ALLOWED);
		break;
	case NVKM_MEM_TARGET_HOST:
		*data |= NVDEF(NV_MMU, VER3_PDE, APERTURE, SYSTEM_COHERENT_MEMORY);
		*data |= NVDEF(NV_MMU, VER3_PDE, PCF, VALID_UNCACHED_ATS_ALLOWED);
		break;
	case NVKM_MEM_TARGET_NCOH:
		*data |= NVDEF(NV_MMU, VER3_PDE, APERTURE, SYSTEM_NON_COHERENT_MEMORY);
		*data |= NVDEF(NV_MMU, VER3_PDE, PCF, VALID_CACHED_ATS_ALLOWED);
		break;
	default:
		WARN_ON(1);
		return false;
	}

	*data |= pt->addr;
	return true;
}

static void
gh100_vmm_pd0_pde(struct nvkm_vmm *vmm, struct nvkm_vmm_pt *pgd, u32 pdei)
{
	struct nvkm_vmm_pt *pgt = pgd->pde[pdei];
	struct nvkm_mmu_pt *pd = pgd->pt[0];
	u64 data[2] = {};

	if (pgt->pt[0] && !gh100_vmm_pde(pgt->pt[0], &data[0]))
		return;
	if (pgt->pt[1] && !gh100_vmm_pde(pgt->pt[1], &data[1]))
		return;

	nvkm_kmap(pd->memory);
	VMM_WO128(pd, vmm, pdei * NV_MMU_VER3_DUAL_PDE__SIZE, data[0], data[1]);
	nvkm_done(pd->memory);
}

static void
gh100_vmm_pd0_sparse(struct nvkm_vmm *vmm,
		     struct nvkm_mmu_pt *pt, u32 pdei, u32 pdes)
{
	const u64 data = NVDEF(NV_MMU, VER3_DUAL_PDE, PCF_BIG, SPARSE_ATS_ALLOWED);

	VMM_FO128(pt, vmm, pdei * NV_MMU_VER3_DUAL_PDE__SIZE, data, 0ULL, pdes);
}

static void
gh100_vmm_pd0_unmap(struct nvkm_vmm *vmm,
		    struct nvkm_mmu_pt *pt, u32 pdei, u32 pdes)
{
	VMM_FO128(pt, vmm, pdei * NV_MMU_VER3_DUAL_PDE__SIZE, 0ULL, 0ULL, pdes);
}

static const struct nvkm_vmm_desc_func
gh100_vmm_desc_pd0 = {
	.unmap = gh100_vmm_pd0_unmap,
	.sparse = gh100_vmm_pd0_sparse,
	.pde = gh100_vmm_pd0_pde,
	.mem = gh100_vmm_pd0_mem,
};

static void
gh100_vmm_pd1_pde(struct nvkm_vmm *vmm, struct nvkm_vmm_pt *pgd, u32 pdei)
{
	struct nvkm_vmm_pt *pgt = pgd->pde[pdei];
	struct nvkm_mmu_pt *pd = pgd->pt[0];
	u64 data = 0;

	if (!gh100_vmm_pde(pgt->pt[0], &data))
		return;

	nvkm_kmap(pd->memory);
	VMM_WO064(pd, vmm, pdei * NV_MMU_VER3_PDE__SIZE, data);
	nvkm_done(pd->memory);
}

static const struct nvkm_vmm_desc_func
gh100_vmm_desc_pd1 = {
	.unmap = gf100_vmm_pgt_unmap,
	.sparse = gh100_vmm_pgt_sparse,
	.pde = gh100_vmm_pd1_pde,
};

static const struct nvkm_vmm_desc
gh100_vmm_desc_16[] = {
	{ LPT, 5,  8, 0x0100, &gh100_vmm_desc_lpt },
	{ PGD, 8, 16, 0x1000, &gh100_vmm_desc_pd0 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 1,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{}
};

static const struct nvkm_vmm_desc
gh100_vmm_desc_12[] = {
	{ SPT, 9,  8, 0x1000, &gh100_vmm_desc_spt },
	{ PGD, 8, 16, 0x1000, &gh100_vmm_desc_pd0 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 9,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{ PGD, 1,  8, 0x1000, &gh100_vmm_desc_pd1 },
	{}
};

static int
gh100_vmm_valid(struct nvkm_vmm *vmm, bool ro, bool priv, u8 kind, u8 comp,
		struct nvkm_vmm_map *map)
{
	const enum nvkm_memory_target target = nvkm_memory_target(map->memory);
	const bool vol = target == NVKM_MEM_TARGET_HOST;
	const struct nvkm_vmm_page *page = map->page;
	u8 kind_inv, pcf;
	int kindn, aper;
	const u8 *kindm;

	map->next = 1ULL << page->shift;
	map->type = 0;

	aper = vmm->func->aper(target);
	if (WARN_ON(aper < 0))
		return aper;

	kindm = vmm->mmu->func->kind(vmm->mmu, &kindn, &kind_inv);
	if (kind >= kindn || kindm[kind] == kind_inv) {
		VMM_DEBUG(vmm, "kind %02x", kind);
		return -EINVAL;
	}

	if (priv) {
		if (ro) {
			if (vol)
				pcf = NV_MMU_VER3_PTE_PCF_PRIVILEGE_RO_ATOMIC_UNCACHED_ACD;
			else
				pcf = NV_MMU_VER3_PTE_PCF_PRIVILEGE_RO_ATOMIC_CACHED_ACD;
		} else {
			if (vol)
				pcf = NV_MMU_VER3_PTE_PCF_PRIVILEGE_RW_ATOMIC_UNCACHED_ACD;
			else
				pcf = NV_MMU_VER3_PTE_PCF_PRIVILEGE_RW_ATOMIC_CACHED_ACD;
		}
	} else {
		if (ro) {
			if (vol)
				pcf = NV_MMU_VER3_PTE_PCF_REGULAR_RO_ATOMIC_UNCACHED_ACD;
			else
				pcf = NV_MMU_VER3_PTE_PCF_REGULAR_RO_ATOMIC_CACHED_ACD;
		} else {
			if (vol)
				pcf = NV_MMU_VER3_PTE_PCF_REGULAR_RW_ATOMIC_UNCACHED_ACD;
			else
				pcf = NV_MMU_VER3_PTE_PCF_REGULAR_RW_ATOMIC_CACHED_ACD;
		}
	}

	map->type |= NVDEF(NV_MMU, VER3_PTE, VALID, TRUE);
	map->type |= NVVAL(NV_MMU, VER3_PTE, APERTURE, aper);
	map->type |= NVVAL(NV_MMU, VER3_PTE, PCF, pcf);
	map->type |= NVVAL(NV_MMU, VER3_PTE, KIND, kind);
	return 0;
}

static const struct nvkm_vmm_func
gh100_vmm = {
	.join = gv100_vmm_join,
	.part = gf100_vmm_part,
	.aper = gf100_vmm_aper,
	.valid = gp100_vmm_valid,
	.valid2 = gh100_vmm_valid,
	.flush = tu102_vmm_flush,
	.page = {
		{ 56, &gh100_vmm_desc_16[5], NVKM_VMM_PAGE_Sxxx },
		{ 47, &gh100_vmm_desc_16[4], NVKM_VMM_PAGE_Sxxx },
		{ 38, &gh100_vmm_desc_16[3], NVKM_VMM_PAGE_Sxxx },
		{ 29, &gh100_vmm_desc_16[2], NVKM_VMM_PAGE_SVxC },
		{ 21, &gh100_vmm_desc_16[1], NVKM_VMM_PAGE_SVxC },
		{ 16, &gh100_vmm_desc_16[0], NVKM_VMM_PAGE_SVxC },
		{ 12, &gh100_vmm_desc_12[0], NVKM_VMM_PAGE_SVHx },
		{}
	}
};

int
gh100_vmm_new(struct nvkm_mmu *mmu, bool managed, u64 addr, u64 size,
	      void *argv, u32 argc, struct lock_class_key *key,
	      const char *name, struct nvkm_vmm **pvmm)
{
	return gp100_vmm_new_(&gh100_vmm, mmu, managed, addr, size,
			      argv, argc, key, name, pvmm);
}