nvml test

nvml.c

#include <stdio.h>
#include <dlfcn.h>

static void *nvml_handle;
typedef struct {
  unsigned long long total;
  unsigned long long free;
  unsigned long long used;
} nvmlMemory_t;

typedef struct {
  char busIdLegacy[16];
  unsigned int domain;
  unsigned int bus;
  unsigned int device;
  unsigned int pciDeviceId;
  // Added in NVML 2.285 API
  unsigned int pciSubSystemId;
  char busId[32];
} nvmlPciInfo_t;


static int (*nvmlInit)(void);
static int (*nvmlDeviceGetMemoryInfo)(void *device, nvmlMemory_t *memory);
static int (*nvmlDeviceGetHandleByIndex) (unsigned int index, void *device);
static int (*nvmlDeviceGetHandleByPciBusId)(char *pciBusId, void* device); 	
static int (*nvmlDeviceGetCount) (unsigned int *count);
static int (*nvmlDeviceGetPciInfo) (void *device, nvmlPciInfo_t *pci);

int main()
{
	if (!nvmlDeviceGetMemoryInfo)
	{
		puts("dlopening libnvidia-ml.so");
		nvml_handle = dlopen("libnvidia-ml.so", RTLD_LAZY);
		if (!nvml_handle)
			nvml_handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY);
		if (!nvml_handle)
		    return 0;

		nvmlInit = dlsym(nvml_handle, "nvmlInit_v2");
		if (!nvmlInit)
			nvmlInit = dlsym(nvml_handle, "nvmlInit");

		nvmlDeviceGetMemoryInfo = dlsym(nvml_handle, "nvmlDeviceGetMemoryInfo");

		nvmlDeviceGetHandleByIndex = dlsym(nvml_handle, "nvmlDeviceGetHandleByIndex_v2");
		if (!nvmlDeviceGetHandleByIndex)
			nvmlDeviceGetHandleByIndex = dlsym(nvml_handle, "nvmlDeviceGetHandleByIndex");

		nvmlDeviceGetHandleByPciBusId = dlsym(nvml_handle, "nvmlDeviceGetHandleByPciBusId_v2");
		if (!nvmlDeviceGetHandleByPciBusId)
			nvmlDeviceGetHandleByPciBusId = dlsym(nvml_handle, "nvmlDeviceGetHandleByPciBusId");

		nvmlDeviceGetCount = dlsym(nvml_handle, "nvmlDeviceGetCount_v2");
		if (!nvmlDeviceGetCount)
			nvmlDeviceGetCount = dlsym(nvml_handle, "nvmlDeviceGetCount");

		nvmlDeviceGetPciInfo = dlsym(nvml_handle, "nvmlDeviceGetPciInfo_v3");
		if (!nvmlDeviceGetPciInfo)
			nvmlDeviceGetPciInfo = dlsym(nvml_handle, "nvmlDeviceGetPciInfo_v2");
		if (!nvmlDeviceGetPciInfo)
			nvmlDeviceGetPciInfo = dlsym(nvml_handle, "nvmlDeviceGetPciInfo");

	}

	puts("init");
	if (nvmlInit())
	{
		puts("nvmlInit error");
		return 1;
	}

	puts("getting count");
	unsigned count;
	if (nvmlDeviceGetCount(&count))
	{
		puts("nvmlDeviceGetCount() error");
		return 1;
	}
	printf("%u devices\n", count);

	for (unsigned i = 0; i < count; ++i)
	{
		void* device;
		if (nvmlDeviceGetHandleByIndex(i, &device))
		{
			puts("nvmlDeviceGetHandleByIndex error");
			continue;
		}

		nvmlMemory_t memory;
		if (nvmlDeviceGetMemoryInfo(device, &memory))
		{
			puts("nvmlDeviceGetMemoryInfo error");
			continue;
		}
		printf("total: %llu\n", memory.total);

		nvmlPciInfo_t pcie;
		if(nvmlDeviceGetPciInfo(device, &pcie))
		{
			puts("nvmlDeviceGetPciInfo error");
			continue;
		}

		printf("bus_id_legacy: %s\n", pcie.busIdLegacy);
		printf("bus_id: %s\n", pcie.busId);
	}



	return 0;
}