/* guest_engines.c - query the guest client RM for the data UVM's register
 * path depends on: virtualization mode, subdevice count, ENGINE LIST,
 * CLASS LIST, CE caps and CE->PCE masks. Uses the NVOS64 serialized alloc
 * path (same as tools/rmq/rmq.c) so device allocs succeed.
 *
 * Build in guest: gcc -O1 -o /tmp/ge guest_engines.c ; run /tmp/ge
 */
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/ioctl.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>

typedef uint32_t NvU32;
typedef uint32_t NvHandle;
typedef uint64_t NvP64;

#define NV_IOCTL_MAGIC      'F'
#define NV_ESC_RM_ALLOC     0x2B
#define NV_ESC_RM_CONTROL   0x2A
#define NV01_ROOT           0x0
#define NV01_ROOT_CLIENT    0x41
#define NV01_DEVICE_0       0x80
#define NV20_SUBDEVICE_0    0x2080

typedef struct {
	NvHandle hRoot, hObjectParent, hObjectNew;
	NvU32 hClass;
	NvP64 pAllocParms;
	NvU32 paramsSize, status;
} NVOS21_PARAMETERS;

typedef struct {
	NvHandle hRoot, hObjectParent, hObjectNew;
	NvU32 hClass;
	NvP64 pAllocParms;
	NvP64 pRightsRequested;
	NvU32 paramsSize, flags, status;
} NVOS64_PARAMETERS;

typedef struct {
	NvHandle hClient, hObject;
	NvU32 cmd, flags;
	NvP64 params;
	NvU32 paramsSize, status;
} NVOS54_PARAMETERS;

static int fd;

static NvU32 rm_alloc(NvHandle hRoot, NvHandle hParent, NvU32 hClass,
		      NvHandle *pOut)
{
	NVOS21_PARAMETERS p;
	memset(&p, 0, sizeof(p));
	p.hRoot = hRoot; p.hObjectParent = hParent; p.hClass = hClass;
	if (ioctl(fd, _IOWR(NV_IOCTL_MAGIC, NV_ESC_RM_ALLOC, NVOS21_PARAMETERS), &p) < 0) {
		perror("alloc ioctl"); return 0xffffffff;
	}
	*pOut = p.hObjectNew;
	return p.status;
}

static NvU32 rm_alloc64(NvHandle hRoot, NvHandle hParent, NvU32 hClass,
			NvHandle *pOut)
{
	NVOS64_PARAMETERS p;
	memset(&p, 0, sizeof(p));
	p.hRoot = hRoot; p.hObjectParent = hParent; p.hClass = hClass;
	p.flags = 1; /* FINN_SERIALIZED */
	if (ioctl(fd, _IOWR(NV_IOCTL_MAGIC, NV_ESC_RM_ALLOC, NVOS64_PARAMETERS), &p) < 0) {
		perror("alloc64 ioctl"); return 0xffffffff;
	}
	*pOut = p.hObjectNew;
	return p.status;
}

static NvU32 rm_control(NvHandle hClient, NvHandle hObject, NvU32 cmd,
			void *params, NvU32 size)
{
	NVOS54_PARAMETERS p;
	memset(&p, 0, sizeof(p));
	p.hClient = hClient; p.hObject = hObject; p.cmd = cmd;
	p.params = (NvP64)params; p.paramsSize = size;
	if (ioctl(fd, _IOWR(NV_IOCTL_MAGIC, NV_ESC_RM_CONTROL, NVOS54_PARAMETERS), &p) < 0) {
		perror("control ioctl"); return 0xffffffff;
	}
	return p.status;
}

#define MAX_CLASSES 256

int main(void)
{
	NvHandle hClient = 0, hDevice = 0, hSubdevice = 0;
	NvU32 st;

	fd = open("/dev/nvidiactl", O_RDWR | O_CLOEXEC);
	if (fd < 0) { perror("open nvidiactl"); return 1; }

	st = rm_alloc(0, 0, getenv("RC") ? NV01_ROOT_CLIENT : NV01_ROOT, &hClient);
	printf("alloc client: st=0x%x hClient=0x%x\n", st, hClient);

	st = rm_alloc64(hClient, hClient, NV01_DEVICE_0, &hDevice);
	printf("alloc device: st=0x%x hDevice=0x%x\n", st, hDevice);
	if (st) return 2;

	st = rm_alloc64(hClient, hDevice, NV20_SUBDEVICE_0, &hSubdevice);
	printf("alloc subdevice: st=0x%x hSub=0x%x\n", st, hSubdevice);

	{
		struct { NvU32 mode; NvU32 isGridBuild; } v = {0xdead, 0xdead};
		st = rm_control(hClient, hDevice, 0x800289, &v, sizeof(v));
		printf("GET_VIRTUALIZATION_MODE: st=0x%x mode=%u (0=NONE,1=NMOS,2=VGX) gridBuild=%u\n",
		       st, v.mode, v.isGridBuild);
	}
	{
		struct { NvU32 n; } s = {0xdead};
		st = rm_control(hClient, hDevice, 0x800280, &s, sizeof(s));
		printf("GET_NUM_SUBDEVICES: st=0x%x n=%u\n", st, s.n);
	}
	{
		struct { NvU32 numClasses; NvU32 classList[MAX_CLASSES]; } c;
		int i;
		memset(&c, 0, sizeof(c));
		st = rm_control(hClient, hDevice, 0x800292, &c, sizeof(c));
		printf("GET_CLASSLIST_V2: st=0x%x numClasses=%u\n", st, c.numClasses);
		if (!st && c.numClasses) {
			printf("  classes:");
			for (i = 0; i < (int)c.numClasses && i < MAX_CLASSES; i++)
				printf(" 0x%x", c.classList[i]);
			printf("\n");
		}
	}
	if (hSubdevice) {
		struct { NvU32 engineCount; NvU32 engineList[0x54]; } e;
		int i;
		memset(&e, 0, sizeof(e));
		st = rm_control(hClient, hSubdevice, 0x20800170, &e, sizeof(e));
		printf("GPU_GET_ENGINES_V2: st=0x%x engineCount=%u\n", st, e.engineCount);
		if (!st && e.engineCount) {
			printf("  engines:");
			for (i = 0; i < (int)e.engineCount && i < 0x54; i++)
				printf(" 0x%x", e.engineList[i]);
			printf("\n");
		}
		for (i = 0; i < (int)e.engineCount && i < 0x54; i++) {
			struct { NvU32 t, mask; } m;
			struct { NvU32 t; unsigned char caps[2]; } cc;
			NvU32 t = e.engineList[i];
			memset(&m, 0, sizeof(m));
			st = rm_control(hClient, hSubdevice, 0x20802a02, &m, sizeof(m));
			printf("  CE_GET_CE_PCE_MASK(type=0x%x): st=0x%x mask=0x%x\n", t, st, m.mask);
			memset(&cc, 0, sizeof(cc));
			cc.t = t;
			st = rm_control(hClient, hSubdevice, 0x20802a03, &cc, sizeof(cc));
			printf("  CE_GET_CAPS_V2(type=0x%x): st=0x%x caps=%02x %02x\n", t, st, cc.caps[0], cc.caps[1]);
		}
	}
	if (hSubdevice) {
		struct {
			NvU32 engineData[16];
			NvU32 pbdmaIds[2];
			NvU32 pbdmaFaultIds[2];
			NvU32 numPbdmas;
			char  engineName[16];
		} e[32];
		struct {
			NvU32 baseIndex, numEntries, bMore;
			typeof(e) entries;
		} t;
		int i, j;
		NvU32 base = 0;
		memset(&t, 0, sizeof(t));
		do {
			t.baseIndex = base;
			st = rm_control(hClient, hSubdevice, 0x20801112, &t, sizeof(t));
			if (st) { printf("FIFO_GET_DEVICE_INFO_TABLE: st=0x%x\n", st); break; }
			for (i = 0; i < (int)t.numEntries; i++) {
				printf("  [%2u] %-8s numPbdmas=%u pbdmaIds=[%u,%u] engineData=[",
				       base + i, t.entries[i].engineName, t.entries[i].numPbdmas,
				       t.entries[i].pbdmaIds[0], t.entries[i].pbdmaIds[1]);
				for (j = 0; j < 16; j++)
					printf("%s%u", j ? "," : "", t.entries[i].engineData[j]);
				printf("]\n");
			}
			base += t.numEntries;
		} while (t.bMore);
	}
	return 0;
}