#include "nnt_device_defs.h"
#include "nnt_gpu.h"

/* Forward declarations */
int is_gb100_pci_device(u_int16_t pci_device_id);
int is_gr100_pci_device(u_int16_t pci_device_id);
int is_gr150_pci_device(u_int16_t pci_device_id);
int is_gpu_pci_device(u_int16_t pci_device_id);

typedef struct pci_id_range_st
{
    unsigned int lower_bound;
    unsigned int upper_bound;
} pci_id_range;

pci_id_range GB100_PCI_IDS[] = {{0x2900, 0x29FF}, {0x2B00, 0x2B7F}, {0x3180, 0x31FF},
                                {0x3200, 0x327F}, {0x3300, 0x33FF}, {0x3400, 0x347F}};

pci_id_range GR100_PCI_IDS[] = {{0x3000, 0x30FF}};

pci_id_range GR150_PCI_IDS[] = {{0x3480, 0x34BF}, {0x34C0, 0x34FF}, {0x3500, 0x353F}, {0x3540, 0x357F}};

int is_gb100_pci_device(u_int16_t pci_device_id)
{
    unsigned int i = 0;
    for (i = 0; i < sizeof(GB100_PCI_IDS) / sizeof(pci_id_range); i++)
    {
        if ((pci_device_id >= GB100_PCI_IDS[i].lower_bound) && (pci_device_id <= GB100_PCI_IDS[i].upper_bound))
        {
            return 1;
        }
    }
    return 0;
}

int is_gr100_pci_device(u_int16_t pci_device_id)
{
    unsigned int i = 0;
    for (i = 0; i < sizeof(GR100_PCI_IDS) / sizeof(pci_id_range); i++)
    {
        if ((pci_device_id >= GR100_PCI_IDS[i].lower_bound) && (pci_device_id <= GR100_PCI_IDS[i].upper_bound))
        {
            return 1;
        }
    }
    return 0;
}

int is_gr150_pci_device(u_int16_t pci_device_id)
{
    unsigned int i = 0;
    for (i = 0; i < sizeof(GR150_PCI_IDS) / sizeof(pci_id_range); i++)
    {
        if ((pci_device_id >= GR150_PCI_IDS[i].lower_bound) && (pci_device_id <= GR150_PCI_IDS[i].upper_bound))
        {
            return 1;
        }
    }
    return 0;
}

int is_gpu_pci_device(u_int16_t pci_device_id)
{
    return (is_gb100_pci_device(pci_device_id) || is_gr100_pci_device(pci_device_id) ||
            is_gr150_pci_device(pci_device_id));
}