/* rmq_probe.c - minimal legacy-ioctl RM client to read device controls:
 *   0x800289 NV0080_CTRL_CMD_GPU_GET_VIRTUALIZATION_MODE
 *   0x800280 NV0080_CTRL_CMD_GPU_GET_NUM_SUBDEVICES
 *   0x800292 NV0080_CTRL_CMD_GPU_GET_CLASSLIST_V2 (count only)
 * plus (optional) the A080 VGPU config query used by UVM's queryVirtMode.
 *
 * Build in guest: gcc -O1 -o /tmp/rmqp rmq_probe.c ; run /tmp/rmqp
 */
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/ioctl.h>
#include <unistd.h>

#define NV_IOCTL_MAGIC 'F'
#define IOCTL_DIR(nr, sz) _IOC(_IOC_READ | _IOC_WRITE, NV_IOCTL_MAGIC, nr, sz)

#define NV_ESC_RM_ALLOC   0x2b
#define NV_ESC_RM_CONTROL 0x2a
#define NV_ESC_RM_FREE    0x29

struct nvos21 {
	uint32_t hRoot, hObjectParent, hObjectNew, hClass;
	uint64_t pAllocParms;
	uint32_t paramsSize, status;
};

struct nvos54 {
	uint32_t hClient, hObject, cmd, flags;
	uint64_t params;
	uint32_t paramsSize, status;
};

#define NV01_ROOT_CLIENT 0x00000041
#define NV01_DEVICE_0    0x00000080

static int fd;

static int rm_alloc(uint32_t hRoot, uint32_t hParent, uint32_t *hNew,
		    uint32_t hClass, void *parms, uint32_t parmsSize)
{
	struct nvos21 a;

	memset(&a, 0, sizeof(a));
	a.hRoot = hRoot;
	a.hObjectParent = hParent;
	a.hObjectNew = *hNew;
	a.hClass = hClass;
	a.pAllocParms = (uint64_t)(uintptr_t)parms;
	a.paramsSize = parmsSize;
	if (ioctl(fd, IOCTL_DIR(NV_ESC_RM_ALLOC, sizeof(a)), &a) < 0) {
		perror("ioctl alloc");
		return -1;
	}
	*hNew = a.hObjectNew;
	return (int)a.status;
}

static int rm_control(uint32_t hClient, uint32_t hObject, uint32_t cmd,
		      void *params, uint32_t paramsSize)
{
	struct nvos54 c;

	memset(&c, 0, sizeof(c));
	c.hClient = hClient;
	c.hObject = hObject;
	c.cmd = cmd;
	c.params = (uint64_t)(uintptr_t)params;
	c.paramsSize = paramsSize;
	if (ioctl(fd, IOCTL_DIR(NV_ESC_RM_CONTROL, sizeof(c)), &c) < 0) {
		perror("ioctl control");
		return -1;
	}
	return (int)c.status;
}

int main(void)
{
	uint32_t hClient = 0, hDevice = 0;
	struct {
		uint32_t deviceId;
		uint32_t hClientShare;
		uint32_t hTargetClient;
		uint32_t hTargetDevice;
		uint32_t flags;
		uint32_t pad0;
		uint64_t vaSpaceSize;
		uint64_t vaStartInternal;
		uint64_t vaLimitInternal;
		uint32_t vaMode;
		uint32_t pad1;
	} devAlloc;
	int st;
	struct {
		uint32_t virtualizationMode;
		uint32_t isGridBuild;
	} virt = { 0xdeadbeef, 0xdeadbeef };
	struct {
		uint32_t numSubDevices;
	} subdev = { 0xdeadbeef };

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

	st = rm_alloc(0, 0, &hClient, NV01_ROOT_CLIENT, NULL, 0);
	printf("alloc client: st=0x%x hClient=0x%x\n", st, hClient);
	if (st != 0)
		return 2;

	memset(&devAlloc, 0, sizeof(devAlloc));
	devAlloc.deviceId = 0;
	st = rm_alloc(hClient, hClient, &hDevice, NV01_DEVICE_0, &devAlloc,
		      sizeof(devAlloc));
	printf("alloc device: st=0x%x hDevice=0x%x\n", st, hDevice);
	if (st != 0)
		return 3;

	st = rm_control(hClient, hDevice, 0x800289, &virt, sizeof(virt));
	printf("GET_VIRTUALIZATION_MODE: st=0x%x mode=0x%x isGridBuild=%u\n",
	       st, virt.virtualizationMode, virt.isGridBuild);

	st = rm_control(hClient, hDevice, 0x800280, &subdev, sizeof(subdev));
	printf("GET_NUM_SUBDEVICES: st=0x%x numSubDevices=%u\n",
	       st, subdev.numSubDevices);

	return 0;
}