/*
 * Hardware Architecture Abstraction Layer
 * Defines hardware-specific configurations for different GPU architectures
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/string.h>

#include "architecture.h"

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Open-IOV");
MODULE_DESCRIPTION("Hardware Architecture Abstraction");

/* Turing architecture (TU102, TU104, TU106) */
static struct hw_architecture arch_turing = {
    .name = "turing",
    .pci_vendor_id = 0x10DE,
    .pci_device_ids = (u32[]){0x1E82, 0x1E83, 0x1E84, 0x1E85, 0x1E86, 0x1E87, 0x1E88, 0x1E89, 0x1E8A, 0x1E8B, 0x1E8C, 0x1E8D, 0x1E8E, 0x1E8F, 0x1ED0, 0x1ED1, 0x1ED2, 0x1ED3, 0x1ED4, 0x1ED5, 0x1ED6, 0x1ED7, 0x1ED8, 0x1ED9, 0x1EDA, 0x1EDB, 0x1EDC, 0x1EDD, 0x1EDE, 0x1EDF, 0},
    .num_device_ids = 30,
    .bar_sizes = {
        .bar0 = 16 * 1024 * 1024,      /* 16MB */
        .bar1 = 256 * 1024 * 1024,     /* 256MB */
        .bar2 = 32 * 1024 * 1024,      /* 32MB */
        .bar3 = 32 * 1024 * 1024,      /* 32MB */
    },
    .doorbell_offset = 0x110C00,
    .shared_mem_offset = 0x1000,
    .shared_mem_size = 4096,
    .msix_table_offset = 0x2000,
    .msix_pba_offset = 0x3000,
    .msix_vectors = 32,
    .max_vfs = 8,
    .vram_size = 16 * 1024 * 1024 * 1024ULL, /* 16GB */
    .page_size = 4096,
    .supports_mig = false,
    .supports_sriov = true,
    .gsp_version = "570.144W",
};

/* Ampere architecture (GA100, GA102, GA104) */
static struct hw_architecture arch_ampere = {
    .name = "ampere",
    .pci_vendor_id = 0x10DE,
    .pci_device_ids = (u32[]){0x20B0, 0x20B1, 0x20B2, 0x20B3, 0x20B4, 0x20B5, 0x20B6, 0x20B7, 0x20B8, 0x20B9, 0x20BA, 0x20BB, 0x20BC, 0x20BD, 0x20BE, 0x20BF, 0x20C0, 0x20C1, 0x20C2, 0x20C3, 0x20C4, 0x20C5, 0x20C6, 0x20C7, 0x20C8, 0x20C9, 0x20CA, 0x20CB, 0x20CC, 0x20CD, 0x20CE, 0x20CF, 0x20D0, 0x20D1, 0x20D2, 0x20D3, 0x20D4, 0x20D5, 0x20D6, 0x20D7, 0x20D8, 0x20D9, 0x20DA, 0x20DB, 0x20DC, 0x20DD, 0x20DE, 0x20DF, 0x20E0, 0x20E1, 0x20E2, 0x20E3, 0x20E4, 0x20E5, 0x20E6, 0x20E7, 0x20E8, 0x20E9, 0x20EA, 0x20EB, 0x20EC, 0x20ED, 0x20EE, 0x20EF, 0},
    .num_device_ids = 96,
    .bar_sizes = {
        .bar0 = 16 * 1024 * 1024,      /* 16MB */
        .bar1 = 512 * 1024 * 1024,     /* 512MB */
        .bar2 = 32 * 1024 * 1024,      /* 32MB */
        .bar3 = 32 * 1024 * 1024,      /* 32MB */
    },
    .doorbell_offset = 0x110C00,
    .shared_mem_offset = 0x1000,
    .shared_mem_size = 4096,
    .msix_table_offset = 0x2000,
    .msix_pba_offset = 0x3000,
    .msix_vectors = 32,
    .max_vfs = 16,
    .vram_size = 80 * 1024 * 1024 * 1024ULL, /* 80GB (A100) */
    .page_size = 4096,
    .supports_mig = true,
    .supports_sriov = true,
    .gsp_version = "570.144W",
};

/* Blackwell architecture (GB100, GB202, GB203) */
static struct hw_architecture arch_blackwell = {
    .name = "blackwell",
    .pci_vendor_id = 0x10DE,
    .pci_device_ids = (u32[]){0x2300, 0x2301, 0x2302, 0x2303, 0x2304, 0x2305, 0x2306, 0x2307, 0x2308, 0x2309, 0x230A, 0x230B, 0x230C, 0x230D, 0x230E, 0x230F, 0},
    .num_device_ids = 16,
    .bar_sizes = {
        .bar0 = 16 * 1024 * 1024,      /* 16MB */
        .bar1 = 512 * 1024 * 1024,     /* 512MB */
        .bar2 = 64 * 1024 * 1024,      /* 64MB */
        .bar3 = 64 * 1024 * 1024,      /* 64MB */
    },
    .doorbell_offset = 0x110C00,
    .shared_mem_offset = 0x1000,
    .shared_mem_size = 4096,
    .msix_table_offset = 0x2000,
    .msix_pba_offset = 0x3000,
    .msix_vectors = 64,
    .max_vfs = 32,
    .vram_size = 192 * 1024 * 1024 * 1024ULL, /* 192GB (B200) */
    .page_size = 4096,
    .supports_mig = true,
    .supports_sriov = true,
    .gsp_version = "580.XX",
};

static struct hw_architecture *architectures[] = {
    &arch_turing,
    &arch_ampere,
    &arch_blackwell,
    NULL
};

struct hw_architecture *arch_get_by_name(const char *name)
{
    for (int i = 0; architectures[i]; i++) {
        if (strcmp(architectures[i]->name, name) == 0)
            return architectures[i];
    }
    return NULL;
}

struct hw_architecture *arch_get_by_pci_id(u32 vendor_id, u32 device_id)
{
    for (int i = 0; architectures[i]; i++) {
        if (architectures[i]->pci_vendor_id == vendor_id) {
            for (int j = 0; j < architectures[i]->num_device_ids; j++) {
                if (architectures[i]->pci_device_ids[j] == device_id)
                    return architectures[i];
            }
        }
    }
    return NULL;
}

struct hw_architecture *arch_get_default(void)
{
    return &arch_turing;
}

const struct hw_bar_config *arch_get_bar_config(struct hw_architecture *arch)
{
    return arch ? &arch->bar_sizes : NULL;
}

u64 arch_get_doorbell_offset(struct hw_architecture *arch)
{
    return arch ? arch->doorbell_offset : 0x110C00;
}

u64 arch_get_shared_mem_offset(struct hw_architecture *arch)
{
    return arch ? arch->shared_mem_offset : 0x1000;
}

u32 arch_get_msix_vectors(struct hw_architecture *arch)
{
    return arch ? arch->msix_vectors : 32;
}

bool arch_supports_mig(struct hw_architecture *arch)
{
    return arch ? arch->supports_mig : false;
}

bool arch_supports_sriov(struct hw_architecture *arch)
{
    return arch ? arch->supports_sriov : false;
}

const char *arch_get_gsp_version(struct hw_architecture *arch)
{
    return arch ? arch->gsp_version : "unknown";
}

EXPORT_SYMBOL_GPL(arch_get_by_name);
EXPORT_SYMBOL_GPL(arch_get_by_pci_id);
EXPORT_SYMBOL_GPL(arch_get_default);
EXPORT_SYMBOL_GPL(arch_get_bar_config);
EXPORT_SYMBOL_GPL(arch_get_doorbell_offset);
EXPORT_SYMBOL_GPL(arch_get_shared_mem_offset);
EXPORT_SYMBOL_GPL(arch_get_msix_vectors);
EXPORT_SYMBOL_GPL(arch_supports_mig);
EXPORT_SYMBOL_GPL(arch_supports_sriov);
EXPORT_SYMBOL_GPL(arch_get_gsp_version);

MODULE_DESCRIPTION("Hardware Architecture Abstraction");
MODULE_LICENSE("GPL");